Skip to main content

limit_cli/
agent_bridge.rs

1use crate::error::CliError;
2use crate::system_prompt::SYSTEM_PROMPT;
3use crate::tools::{
4    AstGrepTool, BashTool, BrowserTool, FileEditTool, FileReadTool, FileWriteTool, GitAddTool,
5    GitCloneTool, GitCommitTool, GitDiffTool, GitLogTool, GitPullTool, GitPushTool, GitStatusTool,
6    GrepTool, LspTool, WebFetchTool, WebSearchTool,
7};
8use chrono::Datelike;
9use futures::StreamExt;
10use limit_agent::executor::{ToolCall, ToolExecutor};
11use limit_agent::registry::ToolRegistry;
12use limit_llm::providers::LlmProvider;
13use limit_llm::types::{Message, Role, Tool as LlmTool, ToolCall as LlmToolCall};
14use limit_llm::ProviderFactory;
15use limit_llm::ProviderResponseChunk;
16use limit_llm::TrackingDb;
17use serde_json::json;
18use tokio::sync::mpsc;
19use tokio_util::sync::CancellationToken;
20use tracing::{debug, instrument};
21
22/// Event types for streaming from agent to REPL
23#[derive(Debug, Clone)]
24#[allow(dead_code)]
25pub enum AgentEvent {
26    Thinking {
27        operation_id: u64,
28    },
29    ToolStart {
30        operation_id: u64,
31        name: String,
32        args: serde_json::Value,
33    },
34    ToolComplete {
35        operation_id: u64,
36        name: String,
37        result: String,
38    },
39    ContentChunk {
40        operation_id: u64,
41        chunk: String,
42    },
43    Done {
44        operation_id: u64,
45    },
46    Cancelled {
47        operation_id: u64,
48    },
49    Error {
50        operation_id: u64,
51        message: String,
52    },
53    TokenUsage {
54        operation_id: u64,
55        input_tokens: u64,
56        output_tokens: u64,
57    },
58}
59
60/// Bridge connecting limit-cli REPL to limit-agent executor and limit-llm client
61pub struct AgentBridge {
62    /// LLM client for communicating with LLM providers
63    llm_client: Box<dyn LlmProvider>,
64    /// Tool executor for running tool calls
65    executor: ToolExecutor,
66    /// List of registered tool names
67    tool_names: Vec<&'static str>,
68    /// Configuration loaded from ~/.limit/config.toml
69    config: limit_llm::Config,
70    /// Event sender for streaming events to REPL
71    event_tx: Option<mpsc::UnboundedSender<AgentEvent>>,
72    /// Token usage tracking database
73    tracking_db: TrackingDb,
74    /// Cancellation token for aborting current operation
75    cancellation_token: Option<CancellationToken>,
76    /// Current operation ID for event tracking
77    operation_id: u64,
78}
79
80impl AgentBridge {
81    /// Create a new AgentBridge with the given configuration
82    ///
83    /// # Arguments
84    /// * `config` - LLM configuration (API key, model, etc.)
85    ///
86    /// # Returns
87    /// A new AgentBridge instance or an error if initialization fails
88    pub fn new(config: limit_llm::Config) -> Result<Self, CliError> {
89        let tracking_db = TrackingDb::new().map_err(|e| CliError::ConfigError(e.to_string()))?;
90        Self::with_tracking_db(config, tracking_db)
91    }
92
93    /// Create a new AgentBridge for testing with an in-memory tracking database
94    #[cfg(test)]
95    pub fn new_for_test(config: limit_llm::Config) -> Result<Self, CliError> {
96        let tracking_db =
97            TrackingDb::new_in_memory().map_err(|e| CliError::ConfigError(e.to_string()))?;
98        Self::with_tracking_db(config, tracking_db)
99    }
100
101    /// Create a new AgentBridge with a custom tracking database
102    pub fn with_tracking_db(
103        config: limit_llm::Config,
104        tracking_db: TrackingDb,
105    ) -> Result<Self, CliError> {
106        let llm_client = ProviderFactory::create_provider(&config)
107            .map_err(|e| CliError::ConfigError(e.to_string()))?;
108
109        let mut tool_registry = ToolRegistry::new();
110        Self::register_tools(&mut tool_registry, &config);
111
112        // Create executor (which takes ownership of registry as Arc)
113        let executor = ToolExecutor::new(tool_registry);
114
115        // Generate tool definitions before giving ownership to executor
116        let tool_names = vec![
117            "file_read",
118            "file_write",
119            "file_edit",
120            "bash",
121            "git_status",
122            "git_diff",
123            "git_log",
124            "git_add",
125            "git_commit",
126            "git_push",
127            "git_pull",
128            "git_clone",
129            "grep",
130            "ast_grep",
131            "lsp",
132            "web_search",
133            "web_fetch",
134            "browser",
135        ];
136
137        Ok(Self {
138            llm_client,
139            executor,
140            tool_names,
141            config,
142            event_tx: None,
143            tracking_db,
144            cancellation_token: None,
145            operation_id: 0,
146        })
147    }
148
149    /// Set the event channel sender for streaming events
150    pub fn set_event_tx(&mut self, tx: mpsc::UnboundedSender<AgentEvent>) {
151        self.event_tx = Some(tx);
152    }
153
154    /// Set the cancellation token and operation ID for this operation
155    pub fn set_cancellation_token(&mut self, token: CancellationToken, operation_id: u64) {
156        debug!("set_cancellation_token: operation_id={}", operation_id);
157        self.cancellation_token = Some(token);
158        self.operation_id = operation_id;
159    }
160
161    /// Clear the cancellation token
162    pub fn clear_cancellation_token(&mut self) {
163        self.cancellation_token = None;
164    }
165
166    /// Register all CLI tools into the tool registry
167    fn register_tools(registry: &mut ToolRegistry, config: &limit_llm::Config) {
168        // File tools
169        registry
170            .register(FileReadTool::new())
171            .expect("Failed to register file_read");
172        registry
173            .register(FileWriteTool::new())
174            .expect("Failed to register file_write");
175        registry
176            .register(FileEditTool::new())
177            .expect("Failed to register file_edit");
178
179        // Bash tool
180        registry
181            .register(BashTool::new())
182            .expect("Failed to register bash");
183
184        // Git tools
185        registry
186            .register(GitStatusTool::new())
187            .expect("Failed to register git_status");
188        registry
189            .register(GitDiffTool::new())
190            .expect("Failed to register git_diff");
191        registry
192            .register(GitLogTool::new())
193            .expect("Failed to register git_log");
194        registry
195            .register(GitAddTool::new())
196            .expect("Failed to register git_add");
197        registry
198            .register(GitCommitTool::new())
199            .expect("Failed to register git_commit");
200        registry
201            .register(GitPushTool::new())
202            .expect("Failed to register git_push");
203        registry
204            .register(GitPullTool::new())
205            .expect("Failed to register git_pull");
206        registry
207            .register(GitCloneTool::new())
208            .expect("Failed to register git_clone");
209
210        // Analysis tools
211        registry
212            .register(GrepTool::new())
213            .expect("Failed to register grep");
214        registry
215            .register(AstGrepTool::new())
216            .expect("Failed to register ast_grep");
217        registry
218            .register(LspTool::new())
219            .expect("Failed to register lsp");
220
221        // Web tools
222        registry
223            .register(WebSearchTool::new())
224            .expect("Failed to register web_search");
225        registry
226            .register(WebFetchTool::new())
227            .expect("Failed to register web_fetch");
228
229        // Browser tool with config
230        let browser_config = crate::tools::browser::BrowserConfig::from(&config.browser);
231        registry
232            .register(BrowserTool::with_config(browser_config))
233            .expect("Failed to register browser");
234    }
235
236    /// Process a user message through the LLM and execute any tool calls
237    ///
238    /// # Arguments
239    /// * `user_input` - The user's message to process
240    /// * `messages` - The conversation history (will be updated in place)
241    ///
242    /// # Returns
243    /// The final response from the LLM or an error
244    #[instrument(skip(self, _messages))]
245    pub async fn process_message(
246        &mut self,
247        user_input: &str,
248        _messages: &mut Vec<Message>,
249    ) -> Result<String, CliError> {
250        // Add system message if this is the first message in the conversation
251        // Note: Some providers (z.ai) don't support system role, but OpenAI-compatible APIs generally do
252        if _messages.is_empty() {
253            let system_message = Message {
254                role: Role::System,
255                content: Some(SYSTEM_PROMPT.to_string()),
256                tool_calls: None,
257                tool_call_id: None,
258            };
259            _messages.push(system_message);
260        }
261
262        // Add user message to history
263        let user_message = Message {
264            role: Role::User,
265            content: Some(user_input.to_string()),
266            tool_calls: None,
267            tool_call_id: None,
268        };
269        _messages.push(user_message);
270
271        // Get tool definitions
272        let tool_definitions = self.get_tool_definitions();
273
274        // Main processing loop
275        let mut full_response = String::new();
276        let mut tool_calls: Vec<LlmToolCall> = Vec::new();
277        let max_iterations = self
278            .config
279            .providers
280            .get(&self.config.provider)
281            .map(|p| p.max_iterations)
282            .unwrap_or(100); // Allow enough iterations for complex tasks
283        let mut iteration = 0;
284
285        while max_iterations == 0 || iteration < max_iterations {
286            iteration += 1;
287            debug!("Agent loop iteration {}", iteration);
288
289            // Send thinking event
290            debug!(
291                "Sending Thinking event with operation_id={}",
292                self.operation_id
293            );
294            self.send_event(AgentEvent::Thinking {
295                operation_id: self.operation_id,
296            });
297
298            // Track timing for token usage
299            let request_start = std::time::Instant::now();
300
301            // Call LLM
302            let mut stream = self
303                .llm_client
304                .send(_messages.clone(), tool_definitions.clone())
305                .await
306                .map_err(|e| CliError::ConfigError(e.to_string()))?;
307
308            tool_calls.clear();
309            let mut current_content = String::new();
310            // Track tool calls: (id) -> (name, args)
311            let mut accumulated_calls: std::collections::HashMap<
312                String,
313                (String, serde_json::Value),
314            > = std::collections::HashMap::new();
315
316            // Process stream chunks with cancellation support
317            loop {
318                // Check for cancellation FIRST (before waiting for stream)
319                if let Some(ref token) = self.cancellation_token {
320                    if token.is_cancelled() {
321                        debug!("Operation cancelled by user (pre-stream check)");
322                        self.send_event(AgentEvent::Cancelled {
323                            operation_id: self.operation_id,
324                        });
325                        return Err(CliError::ConfigError(
326                            "Operation cancelled by user".to_string(),
327                        ));
328                    }
329                }
330
331                // Use tokio::select! to check cancellation while waiting for stream
332                // Using cancellation_token.cancelled() for immediate cancellation detection
333                let chunk_result = if let Some(ref token) = self.cancellation_token {
334                    tokio::select! {
335                        chunk = stream.next() => chunk,
336                        _ = token.cancelled() => {
337                            debug!("Operation cancelled via token while waiting for stream");
338                            self.send_event(AgentEvent::Cancelled {
339                                operation_id: self.operation_id,
340                            });
341                            return Err(CliError::ConfigError("Operation cancelled by user".to_string()));
342                        }
343                    }
344                } else {
345                    stream.next().await
346                };
347
348                let Some(chunk_result) = chunk_result else {
349                    // Stream ended
350                    break;
351                };
352
353                match chunk_result {
354                    Ok(ProviderResponseChunk::ContentDelta(text)) => {
355                        current_content.push_str(&text);
356                        debug!(
357                            "ContentDelta: {} chars (total: {})",
358                            text.len(),
359                            current_content.len()
360                        );
361                        self.send_event(AgentEvent::ContentChunk {
362                            operation_id: self.operation_id,
363                            chunk: text,
364                        });
365                    }
366                    Ok(ProviderResponseChunk::ReasoningDelta(_)) => {
367                        // Ignore reasoning chunks for now
368                    }
369                    Ok(ProviderResponseChunk::ToolCallDelta {
370                        id,
371                        name,
372                        arguments,
373                    }) => {
374                        debug!(
375                            "ToolCallDelta: id={}, name={}, args_len={}",
376                            id,
377                            name,
378                            arguments.to_string().len()
379                        );
380                        // Store/merge tool call arguments
381                        accumulated_calls.insert(id.clone(), (name.clone(), arguments.clone()));
382                    }
383                    Ok(ProviderResponseChunk::Done(usage)) => {
384                        // Track token usage
385                        let duration_ms = request_start.elapsed().as_millis() as u64;
386                        let cost =
387                            calculate_cost(self.model(), usage.input_tokens, usage.output_tokens);
388                        let _ = self.tracking_db.track_request(
389                            self.model(),
390                            usage.input_tokens,
391                            usage.output_tokens,
392                            cost,
393                            duration_ms,
394                        );
395                        // Emit token usage event for TUI display
396                        self.send_event(AgentEvent::TokenUsage {
397                            operation_id: self.operation_id,
398                            input_tokens: usage.input_tokens,
399                            output_tokens: usage.output_tokens,
400                        });
401                        break;
402                    }
403                    Err(e) => {
404                        let error_msg = format!("LLM error: {}", e);
405                        self.send_event(AgentEvent::Error {
406                            operation_id: self.operation_id,
407                            message: error_msg.clone(),
408                        });
409                        return Err(CliError::ConfigError(error_msg));
410                    }
411                }
412            }
413
414            // Convert accumulated calls to Vec<ToolCall>
415            tool_calls = accumulated_calls
416                .into_iter()
417                .map(|(id, (name, args))| LlmToolCall {
418                    id,
419                    tool_type: "function".to_string(),
420                    function: limit_llm::types::FunctionCall {
421                        name,
422                        arguments: args.to_string(),
423                    },
424                })
425                .collect();
426
427            // BUG FIX: Don't accumulate content across iterations
428            // Only store content from the current iteration
429            // If there are tool calls, we'll continue the loop and the LLM will see the tool results
430            // If there are NO tool calls, this is the final response
431            full_response = current_content.clone();
432
433            debug!(
434                "After iter {}: content.len()={}, tool_calls={}, response.len()={}",
435                iteration,
436                current_content.len(),
437                tool_calls.len(),
438                full_response.len()
439            );
440
441            // If no tool calls, we're done
442            if tool_calls.is_empty() {
443                debug!("No tool calls, breaking loop after iteration {}", iteration);
444                break;
445            }
446
447            debug!(
448                "Tool calls found (count={}), continuing to iteration {}",
449                tool_calls.len(),
450                iteration + 1
451            );
452
453            // Execute tool calls - add assistant message with tool_calls
454            // Note: Per OpenAI API spec, when tool_calls are present, content should be null
455            let assistant_message = Message {
456                role: Role::Assistant,
457                content: None, // Don't include content when tool_calls are present
458                tool_calls: Some(tool_calls.clone()),
459                tool_call_id: None,
460            };
461            _messages.push(assistant_message);
462
463            // Convert LLM tool calls to executor tool calls
464            let executor_calls: Vec<ToolCall> = tool_calls
465                .iter()
466                .map(|tc| {
467                    let args: serde_json::Value =
468                        serde_json::from_str(&tc.function.arguments).unwrap_or_default();
469                    ToolCall::new(&tc.id, &tc.function.name, args)
470                })
471                .collect();
472
473            // Send ToolStart event for each tool BEFORE execution
474            for tc in &tool_calls {
475                let args: serde_json::Value =
476                    serde_json::from_str(&tc.function.arguments).unwrap_or_default();
477                self.send_event(AgentEvent::ToolStart {
478                    operation_id: self.operation_id,
479                    name: tc.function.name.clone(),
480                    args,
481                });
482            }
483            // Execute tools
484            let results = self.executor.execute_tools(executor_calls).await;
485
486            // Add tool results to messages (OpenAI format: role=tool, tool_call_id, content)
487            for result in results {
488                let tool_call = tool_calls.iter().find(|tc| tc.id == result.call_id);
489                if let Some(tool_call) = tool_call {
490                    let output_json = match &result.output {
491                        Ok(value) => {
492                            serde_json::to_string(value).unwrap_or_else(|_| "{}".to_string())
493                        }
494                        Err(e) => json!({ "error": e.to_string() }).to_string(),
495                    };
496
497                    self.send_event(AgentEvent::ToolComplete {
498                        operation_id: self.operation_id,
499                        name: tool_call.function.name.clone(),
500                        result: output_json.clone(),
501                    });
502
503                    // OpenAI tool result format
504                    let tool_result_message = Message {
505                        role: Role::Tool,
506                        content: Some(output_json),
507                        tool_calls: None,
508                        tool_call_id: Some(result.call_id),
509                    };
510                    _messages.push(tool_result_message);
511                }
512            }
513        }
514
515        // If we hit max iterations, make one final request to get a response (no tools = forced text)
516        // IMPORTANT: Only do this if max_iterations > 0 (0 means unlimited, so we never "hit" the limit)
517        if max_iterations > 0 && iteration >= max_iterations && !_messages.is_empty() {
518            debug!("Making final LLM call after hitting max iterations (forcing text response)");
519
520            // Add constraint message to force text response
521            let constraint_message = Message {
522                role: Role::User,
523                content: Some(
524                    "We've reached the iteration limit. Please provide a summary of:\n\
525                    1. What you've completed so far\n\
526                    2. What remains to be done\n\
527                    3. Recommended next steps for the user to continue"
528                        .to_string(),
529                ),
530                tool_calls: None,
531                tool_call_id: None,
532            };
533            _messages.push(constraint_message);
534
535            // Send with NO tools to force text response
536            let no_tools: Vec<LlmTool> = vec![];
537            let mut stream = self
538                .llm_client
539                .send(_messages.clone(), no_tools)
540                .await
541                .map_err(|e| CliError::ConfigError(e.to_string()))?;
542
543            // BUG FIX: Replace full_response instead of appending
544            full_response.clear();
545            loop {
546                // Check for cancellation FIRST (before waiting for stream)
547                if let Some(ref token) = self.cancellation_token {
548                    if token.is_cancelled() {
549                        debug!("Operation cancelled by user in final loop (pre-stream check)");
550                        self.send_event(AgentEvent::Cancelled {
551                            operation_id: self.operation_id,
552                        });
553                        return Err(CliError::ConfigError(
554                            "Operation cancelled by user".to_string(),
555                        ));
556                    }
557                }
558
559                // Use tokio::select! to check cancellation while waiting for stream
560                // Using cancellation_token.cancelled() for immediate cancellation detection
561                let chunk_result = if let Some(ref token) = self.cancellation_token {
562                    tokio::select! {
563                        chunk = stream.next() => chunk,
564                        _ = token.cancelled() => {
565                            debug!("Operation cancelled via token while waiting for stream");
566                            self.send_event(AgentEvent::Cancelled {
567                                operation_id: self.operation_id,
568                            });
569                            return Err(CliError::ConfigError("Operation cancelled by user".to_string()));
570                        }
571                    }
572                } else {
573                    stream.next().await
574                };
575
576                let Some(chunk_result) = chunk_result else {
577                    // Stream ended
578                    break;
579                };
580
581                match chunk_result {
582                    Ok(ProviderResponseChunk::ContentDelta(text)) => {
583                        full_response.push_str(&text);
584                        self.send_event(AgentEvent::ContentChunk {
585                            operation_id: self.operation_id,
586                            chunk: text,
587                        });
588                    }
589                    Ok(ProviderResponseChunk::Done(_)) => {
590                        break;
591                    }
592                    Err(e) => {
593                        debug!("Error in final LLM call: {}", e);
594                        break;
595                    }
596                    _ => {}
597                }
598            }
599        }
600
601        // IMPORTANT: Add final assistant response to message history for session persistence
602        // This is crucial for session export/share to work correctly
603        // Only add if we have content AND we haven't already added this response
604        if !full_response.is_empty() {
605            // Find the last assistant message and check if it has content
606            // If it has tool_calls but no content, UPDATE it instead of adding a new one
607            // This prevents accumulation of empty assistant messages in the history
608            let last_assistant_idx = _messages.iter().rposition(|m| m.role == Role::Assistant);
609
610            if let Some(idx) = last_assistant_idx {
611                let last_assistant = &mut _messages[idx];
612
613                // If the last assistant message has no content (tool_calls only), update it
614                if last_assistant.content.is_none()
615                    || last_assistant
616                        .content
617                        .as_ref()
618                        .map(|c| c.is_empty())
619                        .unwrap_or(true)
620                {
621                    last_assistant.content = Some(full_response.clone());
622                    debug!("Updated last assistant message with final response content");
623                } else {
624                    // Last assistant already has content, this shouldn't happen normally
625                    // but we add a new message to be safe
626                    debug!("Last assistant already has content, adding new message");
627                    let final_assistant_message = Message {
628                        role: Role::Assistant,
629                        content: Some(full_response.clone()),
630                        tool_calls: None,
631                        tool_call_id: None,
632                    };
633                    _messages.push(final_assistant_message);
634                }
635            } else {
636                // No assistant message found, add a new one
637                debug!("No assistant message found, adding new message");
638                let final_assistant_message = Message {
639                    role: Role::Assistant,
640                    content: Some(full_response.clone()),
641                    tool_calls: None,
642                    tool_call_id: None,
643                };
644                _messages.push(final_assistant_message);
645            }
646        }
647
648        self.send_event(AgentEvent::Done {
649            operation_id: self.operation_id,
650        });
651        Ok(full_response)
652    }
653
654    /// Get tool definitions formatted for the LLM
655    pub fn get_tool_definitions(&self) -> Vec<LlmTool> {
656        self.tool_names
657            .iter()
658            .map(|name| {
659                let (description, parameters) = Self::get_tool_schema(name);
660                LlmTool {
661                    tool_type: "function".to_string(),
662                    function: limit_llm::types::ToolFunction {
663                        name: name.to_string(),
664                        description,
665                        parameters,
666                    },
667                }
668            })
669            .collect()
670    }
671
672    /// Get the schema (description and parameters) for a tool
673    fn get_tool_schema(name: &str) -> (String, serde_json::Value) {
674        match name {
675            "file_read" => (
676                "Read the contents of a file".to_string(),
677                json!({
678                    "type": "object",
679                    "properties": {
680                        "path": {
681                            "type": "string",
682                            "description": "Path to the file to read"
683                        }
684                    },
685                    "required": ["path"]
686                }),
687            ),
688            "file_write" => (
689                "Write content to a file, creating parent directories if needed".to_string(),
690                json!({
691                    "type": "object",
692                    "properties": {
693                        "path": {
694                            "type": "string",
695                            "description": "Path to the file to write"
696                        },
697                        "content": {
698                            "type": "string",
699                            "description": "Content to write to the file"
700                        }
701                    },
702                    "required": ["path", "content"]
703                }),
704            ),
705            "file_edit" => (
706                "Replace text in a file with new text".to_string(),
707                json!({
708                    "type": "object",
709                    "properties": {
710                        "path": {
711                            "type": "string",
712                            "description": "Path to the file to edit"
713                        },
714                        "old_text": {
715                            "type": "string",
716                            "description": "Text to find and replace"
717                        },
718                        "new_text": {
719                            "type": "string",
720                            "description": "New text to replace with"
721                        }
722                    },
723                    "required": ["path", "old_text", "new_text"]
724                }),
725            ),
726            "bash" => (
727                "Execute a bash command in a shell".to_string(),
728                json!({
729                    "type": "object",
730                    "properties": {
731                        "command": {
732                            "type": "string",
733                            "description": "Bash command to execute"
734                        },
735                        "workdir": {
736                            "type": "string",
737                            "description": "Working directory (default: current directory)"
738                        },
739                        "timeout": {
740                            "type": "integer",
741                            "description": "Timeout in seconds (default: 60)"
742                        }
743                    },
744                    "required": ["command"]
745                }),
746            ),
747            "git_status" => (
748                "Get git repository status".to_string(),
749                json!({
750                    "type": "object",
751                    "properties": {},
752                    "required": []
753                }),
754            ),
755            "git_diff" => (
756                "Get git diff".to_string(),
757                json!({
758                    "type": "object",
759                    "properties": {},
760                    "required": []
761                }),
762            ),
763            "git_log" => (
764                "Get git commit log".to_string(),
765                json!({
766                    "type": "object",
767                    "properties": {
768                        "count": {
769                            "type": "integer",
770                            "description": "Number of commits to show (default: 10)"
771                        }
772                    },
773                    "required": []
774                }),
775            ),
776            "git_add" => (
777                "Add files to git staging area".to_string(),
778                json!({
779                    "type": "object",
780                    "properties": {
781                        "files": {
782                            "type": "array",
783                            "items": {"type": "string"},
784                            "description": "List of file paths to add"
785                        }
786                    },
787                    "required": ["files"]
788                }),
789            ),
790            "git_commit" => (
791                "Create a git commit".to_string(),
792                json!({
793                    "type": "object",
794                    "properties": {
795                        "message": {
796                            "type": "string",
797                            "description": "Commit message"
798                        }
799                    },
800                    "required": ["message"]
801                }),
802            ),
803            "git_push" => (
804                "Push commits to remote repository".to_string(),
805                json!({
806                    "type": "object",
807                    "properties": {
808                        "remote": {
809                            "type": "string",
810                            "description": "Remote name (default: origin)"
811                        },
812                        "branch": {
813                            "type": "string",
814                            "description": "Branch name (default: current branch)"
815                        }
816                    },
817                    "required": []
818                }),
819            ),
820            "git_pull" => (
821                "Pull changes from remote repository".to_string(),
822                json!({
823                    "type": "object",
824                    "properties": {
825                        "remote": {
826                            "type": "string",
827                            "description": "Remote name (default: origin)"
828                        },
829                        "branch": {
830                            "type": "string",
831                            "description": "Branch name (default: current branch)"
832                        }
833                    },
834                    "required": []
835                }),
836            ),
837            "git_clone" => (
838                "Clone a git repository".to_string(),
839                json!({
840                    "type": "object",
841                    "properties": {
842                        "url": {
843                            "type": "string",
844                            "description": "Repository URL to clone"
845                        },
846                        "directory": {
847                            "type": "string",
848                            "description": "Directory to clone into (optional)"
849                        }
850                    },
851                    "required": ["url"]
852                }),
853            ),
854            "grep" => (
855                "Search for text patterns in files using regex".to_string(),
856                json!({
857                    "type": "object",
858                    "properties": {
859                        "pattern": {
860                            "type": "string",
861                            "description": "Regex pattern to search for"
862                        },
863                        "path": {
864                            "type": "string",
865                            "description": "Path to search in (default: current directory)"
866                        }
867                    },
868                    "required": ["pattern"]
869                }),
870            ),
871            "ast_grep" => (
872                "Search code using AST patterns (structural code matching)".to_string(),
873                json!({
874                    "type": "object",
875                    "properties": {
876                        "pattern": {
877                            "type": "string",
878                            "description": "AST pattern to match"
879                        },
880                        "language": {
881                            "type": "string",
882                            "description": "Programming language (rust, typescript, python)"
883                        },
884                        "path": {
885                            "type": "string",
886                            "description": "Path to search in (default: current directory)"
887                        }
888                    },
889                    "required": ["pattern", "language"]
890                }),
891            ),
892            "lsp" => (
893                "Perform Language Server Protocol operations (goto_definition, find_references)"
894                    .to_string(),
895                json!({
896                    "type": "object",
897                    "properties": {
898                        "command": {
899                            "type": "string",
900                            "description": "LSP command: goto_definition or find_references"
901                        },
902                        "file_path": {
903                            "type": "string",
904                            "description": "Path to the file"
905                        },
906                        "position": {
907                            "type": "object",
908                            "description": "Position in the file (line, character)",
909                            "properties": {
910                                "line": {"type": "integer"},
911                                "character": {"type": "integer"}
912                            },
913                            "required": ["line", "character"]
914                        }
915                    },
916                    "required": ["command", "file_path", "position"]
917                }),
918            ),
919            "web_search" => (
920                format!("Search the web using Exa AI. Returns results with titles, URLs, and content snippets. Use for current information beyond knowledge cutoff. The current year is {} - use this year when searching for recent information.", chrono::Local::now().year()),
921                json!({
922                    "type": "object",
923                    "properties": {
924                        "query": {
925                            "type": "string",
926                            "description": format!("Search query. Be specific for better results (e.g., 'Rust async tutorial {}' rather than 'Rust')", chrono::Local::now().year())
927                        },
928                        "numResults": {
929                            "type": "integer",
930                            "description": "Number of results to return (default: 8, max: 20)",
931                            "default": 8
932                        }
933                    },
934                    "required": ["query"]
935                }),
936            ),
937            "web_fetch" => (
938                "Fetch content from a URL. Converts HTML to markdown format by default. Use when user provides a URL or after web_search to read full content of a specific result.".to_string(),
939                json!({
940                    "type": "object",
941                    "properties": {
942                        "url": {
943                            "type": "string",
944                            "description": "URL to fetch (must start with http:// or https://)"
945                        },
946                        "format": {
947                            "type": "string",
948                            "enum": ["markdown", "text", "html"],
949                            "default": "markdown",
950                            "description": "Output format (default: markdown)"
951                        }
952                    },
953                    "required": ["url"]
954                }),
955            ),
956            "browser" => (
957                "Browser automation for testing, scraping, and web interaction. Use snapshot-ref workflow: open URL, take snapshot, use refs from snapshot for interactions. Supports Chrome and Lightpanda engines.".to_string(),
958                json!({
959                    "type": "object",
960                    "properties": {
961                        "action": {
962                            "type": "string",
963                            "enum": [
964                                // Core
965                                "open", "close", "snapshot",
966                                // Interaction
967                                "click", "dblclick", "fill", "type", "press", "hover", "select",
968                                "focus", "check", "uncheck", "scrollintoview", "drag", "upload",
969                                // Navigation
970                                "back", "forward", "reload",
971                                // Query
972                                "screenshot", "pdf", "eval", "get", "get_attr", "get_count", "get_box", "get_styles",
973                                "find", "is", "download",
974                                // Waiting
975                                "wait", "wait_for_text", "wait_for_url", "wait_for_load", "wait_for_download", "wait_for_fn", "wait_for_state",
976                                // Tabs & Dialogs
977                                "tab_list", "tab_new", "tab_close", "tab_select", "dialog_accept", "dialog_dismiss",
978                                // Storage & Network
979                                "cookies", "cookies_set", "storage_get", "storage_set", "network_requests",
980                                // Settings
981                                "set_viewport", "set_device", "set_geo",
982                                // State
983                                "scroll"
984                            ],
985                            "description": "Browser action to perform"
986                        },
987                        // Core
988                        "url": {
989                            "type": "string",
990                            "description": "URL to open (required for 'open' action)"
991                        },
992                        // Interaction
993                        "selector": {
994                            "type": "string",
995                            "description": "Element selector or ref (for click, fill, type, hover, select, focus, check, uncheck, scrollintoview, get_attr, get_count, get_box, get_styles, is, download, upload)"
996                        },
997                        "text": {
998                            "type": "string",
999                            "description": "Text to input (for fill, type actions)"
1000                        },
1001                        "key": {
1002                            "type": "string",
1003                            "description": "Key to press (required for 'press' action)"
1004                        },
1005                        "value": {
1006                            "type": "string",
1007                            "description": "Value (for select, cookies_set, storage_set)"
1008                        },
1009                        "target": {
1010                            "type": "string",
1011                            "description": "Target selector (for drag action)"
1012                        },
1013                        "files": {
1014                            "type": "array",
1015                            "items": {"type": "string"},
1016                            "description": "File paths to upload (for upload action)"
1017                        },
1018                        // Query
1019                        "path": {
1020                            "type": "string",
1021                            "description": "File path (for screenshot, pdf, download actions)"
1022                        },
1023                        "script": {
1024                            "type": "string",
1025                            "description": "JavaScript to evaluate (required for 'eval' and 'wait_for_fn' actions)"
1026                        },
1027                        "get_what": {
1028                            "type": "string",
1029                            "enum": ["text", "html", "value", "url", "title"],
1030                            "description": "What to get (required for 'get' action)"
1031                        },
1032                        "attr": {
1033                            "type": "string",
1034                            "description": "Attribute name (for get_attr action)"
1035                        },
1036                        // Find
1037                        "locator_type": {
1038                            "type": "string",
1039                            "enum": ["role", "text", "label", "placeholder", "alt", "title", "testid", "css", "xpath"],
1040                            "description": "Locator strategy (for find action)"
1041                        },
1042                        "locator_value": {
1043                            "type": "string",
1044                            "description": "Locator value (for find action)"
1045                        },
1046                        "find_action": {
1047                            "type": "string",
1048                            "enum": ["click", "fill", "text", "count", "first", "last", "nth", "hover", "focus", "check", "uncheck"],
1049                            "description": "Action to perform on found element (for find action)"
1050                        },
1051                        "action_value": {
1052                            "type": "string",
1053                            "description": "Value for find action (optional)"
1054                        },
1055                        // Waiting
1056                        "wait_for": {
1057                            "type": "string",
1058                            "description": "Wait condition (for wait action)"
1059                        },
1060                        "state": {
1061                            "type": "string",
1062                            "enum": ["visible", "hidden", "attached", "detached", "enabled", "disabled", "networkidle", "domcontentloaded", "load"],
1063                            "description": "State to wait for (for wait_for_state, wait_for_load actions)"
1064                        },
1065                        // State check
1066                        "what": {
1067                            "type": "string",
1068                            "enum": ["visible", "hidden", "enabled", "disabled", "editable"],
1069                            "description": "State to check (required for 'is' action)"
1070                        },
1071                        // Scroll
1072                        "direction": {
1073                            "type": "string",
1074                            "enum": ["up", "down", "left", "right"],
1075                            "description": "Scroll direction (for scroll action)"
1076                        },
1077                        "pixels": {
1078                            "type": "integer",
1079                            "description": "Pixels to scroll (optional for scroll action)"
1080                        },
1081                        // Tabs
1082                        "index": {
1083                            "type": "integer",
1084                            "description": "Tab index (for tab_close, tab_select actions)"
1085                        },
1086                        // Dialogs
1087                        "dialog_text": {
1088                            "type": "string",
1089                            "description": "Text for prompt dialog (for dialog_accept action)"
1090                        },
1091                        // Storage
1092                        "storage_type": {
1093                            "type": "string",
1094                            "enum": ["local", "session"],
1095                            "description": "Storage type (for storage_get, storage_set actions)"
1096                        },
1097                        "key_name": {
1098                            "type": "string",
1099                            "description": "Storage key name (for storage_get, storage_set actions)"
1100                        },
1101                        // Network
1102                        "filter": {
1103                            "type": "string",
1104                            "description": "Network request filter (optional for network_requests action)"
1105                        },
1106                        // Settings
1107                        "width": {
1108                            "type": "integer",
1109                            "description": "Viewport width (for set_viewport action)"
1110                        },
1111                        "height": {
1112                            "type": "integer",
1113                            "description": "Viewport height (for set_viewport action)"
1114                        },
1115                        "scale": {
1116                            "type": "number",
1117                            "description": "Device scale factor (optional for set_viewport action)"
1118                        },
1119                        "device_name": {
1120                            "type": "string",
1121                            "description": "Device name to emulate (for set_device action)"
1122                        },
1123                        "latitude": {
1124                            "type": "number",
1125                            "description": "Latitude (for set_geo action)"
1126                        },
1127                        "longitude": {
1128                            "type": "number",
1129                            "description": "Longitude (for set_geo action)"
1130                        },
1131                        // Cookie
1132                        "name": {
1133                            "type": "string",
1134                            "description": "Cookie name (for cookies_set action)"
1135                        },
1136                        // Engine
1137                        "engine": {
1138                            "type": "string",
1139                            "enum": ["chrome", "lightpanda"],
1140                            "default": "chrome",
1141                            "description": "Browser engine to use"
1142                        }
1143                    },
1144                    "required": ["action"]
1145                }),
1146            ),
1147            _ => (
1148                format!("Tool: {}", name),
1149                json!({
1150                    "type": "object",
1151                    "properties": {},
1152                    "required": []
1153                }),
1154            ),
1155        }
1156    }
1157
1158    /// Send an event through the event channel
1159    fn send_event(&self, event: AgentEvent) {
1160        if let Some(ref tx) = self.event_tx {
1161            let _ = tx.send(event);
1162        }
1163    }
1164
1165    /// Check if the bridge is ready to process messages
1166    #[allow(dead_code)]
1167    pub fn is_ready(&self) -> bool {
1168        self.config
1169            .providers
1170            .get(&self.config.provider)
1171            .map(|p| p.api_key_or_env(&self.config.provider).is_some())
1172            .unwrap_or(false)
1173    }
1174
1175    /// Get the current model name
1176    pub fn model(&self) -> &str {
1177        self.config
1178            .providers
1179            .get(&self.config.provider)
1180            .map(|p| p.model.as_str())
1181            .unwrap_or("")
1182    }
1183
1184    /// Get the max tokens setting
1185    pub fn max_tokens(&self) -> u32 {
1186        self.config
1187            .providers
1188            .get(&self.config.provider)
1189            .map(|p| p.max_tokens)
1190            .unwrap_or(4096)
1191    }
1192
1193    /// Get the timeout setting
1194    pub fn timeout(&self) -> u64 {
1195        self.config
1196            .providers
1197            .get(&self.config.provider)
1198            .map(|p| p.timeout)
1199            .unwrap_or(60)
1200    }
1201}
1202/// Calculate cost based on model pricing (per 1M tokens)
1203fn calculate_cost(model: &str, input_tokens: u64, output_tokens: u64) -> f64 {
1204    let (input_price, output_price) = match model {
1205        // Claude 3.5 Sonnet: $3/1M input, $15/1M output
1206        "claude-3-5-sonnet-20241022" | "claude-3-5-sonnet" => (3.0, 15.0),
1207        // GPT-4: $30/1M input, $60/1M output
1208        "gpt-4" => (30.0, 60.0),
1209        // GPT-4 Turbo: $10/1M input, $30/1M output
1210        "gpt-4-turbo" | "gpt-4-turbo-preview" => (10.0, 30.0),
1211        // Default: no cost tracking
1212        _ => (0.0, 0.0),
1213    };
1214    (input_tokens as f64 * input_price / 1_000_000.0)
1215        + (output_tokens as f64 * output_price / 1_000_000.0)
1216}
1217
1218#[cfg(test)]
1219mod tests {
1220    use super::*;
1221    use limit_llm::{BrowserConfigSection, Config as LlmConfig, ProviderConfig};
1222    use std::collections::HashMap;
1223
1224    #[tokio::test]
1225    async fn test_agent_bridge_new() {
1226        let mut providers = HashMap::new();
1227        providers.insert(
1228            "anthropic".to_string(),
1229            ProviderConfig {
1230                api_key: Some("test-key".to_string()),
1231                model: "claude-3-5-sonnet-20241022".to_string(),
1232                base_url: None,
1233                max_tokens: 4096,
1234                timeout: 60,
1235                max_iterations: 100,
1236                thinking_enabled: false,
1237                clear_thinking: true,
1238            },
1239        );
1240        let config = LlmConfig {
1241            provider: "anthropic".to_string(),
1242            providers,
1243            browser: BrowserConfigSection::default(),
1244        };
1245
1246        let bridge = AgentBridge::new(config).unwrap();
1247        assert!(bridge.is_ready());
1248    }
1249
1250    #[tokio::test]
1251    async fn test_agent_bridge_new_no_api_key() {
1252        let mut providers = HashMap::new();
1253        providers.insert(
1254            "anthropic".to_string(),
1255            ProviderConfig {
1256                api_key: None,
1257                model: "claude-3-5-sonnet-20241022".to_string(),
1258                base_url: None,
1259                max_tokens: 4096,
1260                timeout: 60,
1261                max_iterations: 100,
1262                thinking_enabled: false,
1263                clear_thinking: true,
1264            },
1265        );
1266        let config = LlmConfig {
1267            provider: "anthropic".to_string(),
1268            providers,
1269            browser: BrowserConfigSection::default(),
1270        };
1271
1272        let result = AgentBridge::new(config);
1273        assert!(result.is_err());
1274    }
1275
1276    #[tokio::test]
1277    async fn test_get_tool_definitions() {
1278        let mut providers = HashMap::new();
1279        providers.insert(
1280            "anthropic".to_string(),
1281            ProviderConfig {
1282                api_key: Some("test-key".to_string()),
1283                model: "claude-3-5-sonnet-20241022".to_string(),
1284                base_url: None,
1285                max_tokens: 4096,
1286                timeout: 60,
1287                max_iterations: 100,
1288                thinking_enabled: false,
1289                clear_thinking: true,
1290            },
1291        );
1292        let config = LlmConfig {
1293            provider: "anthropic".to_string(),
1294            providers,
1295            browser: BrowserConfigSection::default(),
1296        };
1297
1298        let bridge = AgentBridge::new(config).unwrap();
1299        let definitions = bridge.get_tool_definitions();
1300
1301        assert_eq!(definitions.len(), 18);
1302
1303        // Check file_read tool definition
1304        let file_read = definitions
1305            .iter()
1306            .find(|d| d.function.name == "file_read")
1307            .unwrap();
1308        assert_eq!(file_read.tool_type, "function");
1309        assert_eq!(file_read.function.name, "file_read");
1310        assert!(file_read.function.description.contains("Read"));
1311
1312        // Check bash tool definition
1313        let bash = definitions
1314            .iter()
1315            .find(|d| d.function.name == "bash")
1316            .unwrap();
1317        assert_eq!(bash.function.name, "bash");
1318        assert!(bash.function.parameters["required"]
1319            .as_array()
1320            .unwrap()
1321            .contains(&"command".into()));
1322    }
1323
1324    #[test]
1325    fn test_get_tool_schema() {
1326        let (desc, params) = AgentBridge::get_tool_schema("file_read");
1327        assert!(desc.contains("Read"));
1328        assert_eq!(params["properties"]["path"]["type"], "string");
1329        assert!(params["required"]
1330            .as_array()
1331            .unwrap()
1332            .contains(&"path".into()));
1333
1334        let (desc, params) = AgentBridge::get_tool_schema("bash");
1335        assert!(desc.contains("bash"));
1336        assert_eq!(params["properties"]["command"]["type"], "string");
1337
1338        let (desc, _params) = AgentBridge::get_tool_schema("unknown_tool");
1339        assert!(desc.contains("unknown_tool"));
1340    }
1341
1342    #[test]
1343    fn test_is_ready() {
1344        let mut providers = HashMap::new();
1345        providers.insert(
1346            "anthropic".to_string(),
1347            ProviderConfig {
1348                api_key: Some("test-key".to_string()),
1349                model: "claude-3-5-sonnet-20241022".to_string(),
1350                base_url: None,
1351                max_tokens: 4096,
1352                timeout: 60,
1353                max_iterations: 100,
1354                thinking_enabled: false,
1355                clear_thinking: true,
1356            },
1357        );
1358        let config_with_key = LlmConfig {
1359            provider: "anthropic".to_string(),
1360            providers,
1361            browser: BrowserConfigSection::default(),
1362        };
1363
1364        let bridge = AgentBridge::new(config_with_key).unwrap();
1365        assert!(bridge.is_ready());
1366    }
1367}