Skip to main content

limit_cli/
agent_bridge.rs

1use crate::error::CliError;
2use crate::system_prompt::get_system_prompt;
3use crate::tools::{
4    AstGrepTool, BashTool, BrowserTool, FileEditTool, FileReadTool, FileWriteTool, GitAddTool,
5    GitCloneTool, GitCommitTool, GitDiffTool, GitLogTool, GitPullTool, GitPushTool, GitStatusTool,
6    WebFetchTool, WebSearchTool,
7};
8use chrono::Datelike;
9use futures::StreamExt;
10use limit_agent::executor::{ToolCall, ToolExecutor};
11use limit_agent::registry::ToolRegistry;
12use limit_llm::apply_cache_control;
13use limit_llm::providers::LlmProvider;
14use limit_llm::types::{Message, Role, Tool as LlmTool, ToolCall as LlmToolCall};
15use limit_llm::ModelHandoff;
16use limit_llm::ProviderFactory;
17use limit_llm::ProviderResponseChunk;
18use limit_llm::Summarizer;
19use limit_llm::TrackingDb;
20use serde_json::json;
21use std::cell::RefCell;
22use std::collections::hash_map::DefaultHasher;
23use std::hash::{Hash, Hasher};
24use tokio::sync::mpsc;
25use tokio_util::sync::CancellationToken;
26use tracing::{debug, instrument, trace};
27
28/// Event types for streaming from agent to REPL
29#[derive(Debug, Clone)]
30#[allow(dead_code)]
31pub enum AgentEvent {
32    Thinking {
33        operation_id: u64,
34    },
35    ToolStart {
36        operation_id: u64,
37        name: String,
38        args: serde_json::Value,
39    },
40    ToolComplete {
41        operation_id: u64,
42        name: String,
43        result: String,
44    },
45    ResponseStart {
46        operation_id: u64,
47    },
48    ContentChunk {
49        operation_id: u64,
50        chunk: String,
51    },
52    Done {
53        operation_id: u64,
54    },
55    Cancelled {
56        operation_id: u64,
57    },
58    Error {
59        operation_id: u64,
60        message: String,
61    },
62    TokenUsage {
63        operation_id: u64,
64        input_tokens: u64,
65        output_tokens: u64,
66    },
67}
68
69/// Result from processing a message
70#[derive(Debug, Clone, Default)]
71pub struct ProcessResult {
72    pub response: String,
73    pub input_tokens: u64,
74    pub output_tokens: u64,
75}
76
77/// Maximum number of recent tool calls to keep for deduplication
78const MAX_RECENT_TOOL_CALLS: usize = 20;
79
80/// Maximum characters in tool result before truncation
81const MAX_TOOL_RESULT_CHARS: usize = 10000;
82
83/// Bridge connecting limit-cli REPL to limit-agent executor and limit-llm client
84pub struct AgentBridge {
85    /// LLM client for communicating with LLM providers
86    llm_client: Box<dyn LlmProvider>,
87    /// Tool executor for running tool calls
88    executor: ToolExecutor,
89    /// List of registered tool names
90    tool_names: Vec<&'static str>,
91    /// Configuration loaded from ~/.limit/config.toml
92    config: limit_llm::Config,
93    /// Event sender for streaming events to REPL
94    event_tx: Option<mpsc::UnboundedSender<AgentEvent>>,
95    /// Token usage tracking database
96    tracking_db: TrackingDb,
97    cancellation_token: Option<CancellationToken>,
98    operation_id: u64,
99    recent_tool_calls: RefCell<Vec<(String, u64)>>,
100    handoff: ModelHandoff,
101    summarizer: Option<Summarizer>,
102    last_context_percent: RefCell<usize>,
103}
104
105impl AgentBridge {
106    /// Create a new AgentBridge with the given configuration
107    /// Create a new AgentBridge with the given configuration
108    /// Create a new AgentBridge with the given configuration
109    ///
110    /// # Arguments
111    /// * `config` - LLM configuration (API key, model, etc.)
112    ///
113    /// # Returns
114    /// A new AgentBridge instance or an error if initialization fails
115    pub fn new(config: limit_llm::Config) -> Result<Self, CliError> {
116        let tracking_db = TrackingDb::new().map_err(|e| CliError::ConfigError(e.to_string()))?;
117        Self::with_tracking_db(config, tracking_db)
118    }
119
120    /// Create a new AgentBridge for testing with an in-memory tracking database
121    #[cfg(test)]
122    pub fn new_for_test(config: limit_llm::Config) -> Result<Self, CliError> {
123        let tracking_db =
124            TrackingDb::new_in_memory().map_err(|e| CliError::ConfigError(e.to_string()))?;
125        Self::with_tracking_db(config, tracking_db)
126    }
127
128    /// Create a new AgentBridge with a custom tracking database
129    pub fn with_tracking_db(
130        config: limit_llm::Config,
131        tracking_db: TrackingDb,
132    ) -> Result<Self, CliError> {
133        let llm_client = ProviderFactory::create_provider(&config)
134            .map_err(|e| CliError::ConfigError(e.to_string()))?;
135
136        let mut tool_registry = ToolRegistry::new();
137        Self::register_tools(&mut tool_registry, &config);
138
139        let executor = ToolExecutor::new(tool_registry);
140
141        let tool_names = vec![
142            "file_read",
143            "file_write",
144            "file_edit",
145            "bash",
146            "git_status",
147            "git_diff",
148            "git_log",
149            "git_add",
150            "git_commit",
151            "git_push",
152            "git_pull",
153            "git_clone",
154            // "grep",           // TEMP: disabled for ast_grep testing
155            "ast_grep",
156            // "lsp",            // TEMP: disabled for ast_grep testing
157            "web_search",
158            "web_fetch",
159            "browser",
160        ];
161
162        Ok(Self {
163            llm_client,
164            executor,
165            tool_names,
166            config,
167            event_tx: None,
168            tracking_db,
169            cancellation_token: None,
170            operation_id: 0,
171            recent_tool_calls: RefCell::new(Vec::new()),
172            handoff: ModelHandoff::new(),
173            summarizer: None,
174            last_context_percent: RefCell::new(0),
175        })
176    }
177
178    /// Set the event channel sender for streaming events
179    pub fn set_event_tx(&mut self, tx: mpsc::UnboundedSender<AgentEvent>) {
180        self.event_tx = Some(tx);
181    }
182
183    /// Set the cancellation token and operation ID for this operation
184    pub fn set_cancellation_token(&mut self, token: CancellationToken, operation_id: u64) {
185        debug!("set_cancellation_token: operation_id={}", operation_id);
186        self.cancellation_token = Some(token);
187        self.operation_id = operation_id;
188    }
189
190    /// Clear the cancellation token
191    pub fn clear_cancellation_token(&mut self) {
192        self.cancellation_token = None;
193    }
194
195    async fn maybe_compact(&self, messages: &mut Vec<Message>) {
196        if !self.config.compaction.enabled {
197            return;
198        }
199
200        let context_window: usize = 200_000;
201        let target_tokens = (context_window * 6) / 10;
202        let warn_tokens = context_window / 2;
203        let current_tokens = self.handoff.count_total_tokens(messages);
204        let current_pct = (current_tokens * 100) / context_window;
205
206        if current_tokens > warn_tokens && current_tokens <= target_tokens {
207            let last_pct = *self.last_context_percent.borrow();
208            if current_pct != last_pct {
209                tracing::warn!(
210                    "Context at {}% ({} tokens). Compaction will trigger at 60%.",
211                    current_pct,
212                    current_tokens
213                );
214                *self.last_context_percent.borrow_mut() = current_pct;
215            }
216        } else if current_tokens <= warn_tokens {
217            *self.last_context_percent.borrow_mut() = 0;
218        }
219
220        if current_tokens <= target_tokens {
221            return;
222        }
223
224        let keep_recent = self.config.compaction.keep_recent_tokens as usize;
225
226        if let Some(ref summarizer) = self.summarizer {
227            if let Some(cut_idx) = self.handoff.find_cut_point(messages, keep_recent) {
228                if cut_idx > 0 {
229                    let to_summarize = &messages[..cut_idx];
230
231                    match summarizer.summarize(to_summarize, None).await {
232                        Ok(summary) => {
233                            let summary_msg = Message {
234                                role: Role::User,
235                                content: Some(format!(
236                                    "<context_summary>\n{}\n</context_summary>",
237                                    summary
238                                )),
239                                tool_calls: None,
240                                tool_call_id: None,
241                                cache_control: None,
242                            };
243
244                            let mut new_messages = vec![summary_msg];
245                            new_messages.extend(messages[cut_idx..].to_vec());
246                            *messages = new_messages;
247
248                            debug!(
249                                "Compacted via summarization: {} messages -> {} messages",
250                                cut_idx,
251                                messages.len()
252                            );
253                            return;
254                        }
255                        Err(e) => {
256                            debug!("Summarization failed, falling back to truncation: {}", e);
257                        }
258                    }
259                }
260            }
261        }
262
263        let compacted = self.handoff.compact_messages(messages, target_tokens);
264        *messages = compacted;
265    }
266
267    fn hash_tool_call(tool_name: &str, args: &serde_json::Value) -> u64 {
268        let mut hasher = DefaultHasher::new();
269        tool_name.hash(&mut hasher);
270        args.to_string().hash(&mut hasher);
271        hasher.finish()
272    }
273
274    fn check_duplicate_tool_call(&self, tool_name: &str, args: &serde_json::Value) -> bool {
275        let hash = Self::hash_tool_call(tool_name, args);
276        self.recent_tool_calls
277            .borrow()
278            .iter()
279            .any(|(name, h)| *name == tool_name && *h == hash)
280    }
281
282    fn record_tool_call(&self, tool_name: &str, args: &serde_json::Value) {
283        let hash = Self::hash_tool_call(tool_name, args);
284        self.recent_tool_calls
285            .borrow_mut()
286            .push((tool_name.to_string(), hash));
287        if self.recent_tool_calls.borrow().len() > MAX_RECENT_TOOL_CALLS {
288            self.recent_tool_calls.borrow_mut().remove(0);
289        }
290    }
291
292    /// Register all CLI tools into the tool registry
293    fn register_tools(registry: &mut ToolRegistry, config: &limit_llm::Config) {
294        // File tools
295        registry
296            .register(FileReadTool::new())
297            .expect("Failed to register file_read");
298        registry
299            .register(FileWriteTool::new())
300            .expect("Failed to register file_write");
301        registry
302            .register(FileEditTool::new())
303            .expect("Failed to register file_edit");
304
305        // Bash tool
306        registry
307            .register(BashTool::new())
308            .expect("Failed to register bash");
309
310        // Git tools
311        registry
312            .register(GitStatusTool::new())
313            .expect("Failed to register git_status");
314        registry
315            .register(GitDiffTool::new())
316            .expect("Failed to register git_diff");
317        registry
318            .register(GitLogTool::new())
319            .expect("Failed to register git_log");
320        registry
321            .register(GitAddTool::new())
322            .expect("Failed to register git_add");
323        registry
324            .register(GitCommitTool::new())
325            .expect("Failed to register git_commit");
326        registry
327            .register(GitPushTool::new())
328            .expect("Failed to register git_push");
329        registry
330            .register(GitPullTool::new())
331            .expect("Failed to register git_pull");
332        registry
333            .register(GitCloneTool::new())
334            .expect("Failed to register git_clone");
335
336        // Analysis tools
337        // TEMP: disabled for ast_grep testing
338        // registry
339        //     .register(GrepTool::new())
340        //     .expect("Failed to register grep");
341        registry
342            .register(AstGrepTool::new())
343            .expect("Failed to register ast_grep");
344        // registry
345        //     .register(LspTool::new())
346        //     .expect("Failed to register lsp");
347
348        // Web tools
349        registry
350            .register(WebSearchTool::new())
351            .expect("Failed to register web_search");
352        registry
353            .register(WebFetchTool::new())
354            .expect("Failed to register web_fetch");
355
356        // Browser tool
357        let browser_config = crate::tools::browser::BrowserConfig::from(&config.browser);
358        registry
359            .register(BrowserTool::with_config(browser_config))
360            .expect("Failed to register browser");
361    }
362
363    /// Process a user message through the LLM and execute any tool calls
364    ///
365    /// # Arguments
366    /// * `user_input` - The user's message to process
367    /// * `messages` - The conversation history (will be updated in place)
368    ///
369    /// # Returns
370    /// The final response from the LLM or an error
371    #[instrument(skip(self, _messages, user_input))]
372    pub async fn process_message(
373        &mut self,
374        user_input: &str,
375        _messages: &mut Vec<Message>,
376    ) -> Result<ProcessResult, CliError> {
377        // Add system message if this is the first message in the conversation
378        // Note: Some providers (z.ai) don't support system role, but OpenAI-compatible APIs generally do
379        if _messages.is_empty() {
380            let system_message = Message {
381                role: Role::System,
382                content: Some(get_system_prompt()),
383                tool_calls: None,
384                tool_call_id: None,
385                cache_control: None,
386            };
387            _messages.push(system_message);
388        }
389
390        // Add user message to history
391        let user_message = Message {
392            role: Role::User,
393            content: Some(user_input.to_string()),
394            tool_calls: None,
395            tool_call_id: None,
396            cache_control: None,
397        };
398        _messages.push(user_message);
399
400        // Get tool definitions
401        let tool_definitions = self.get_tool_definitions();
402
403        // Main processing loop
404        let mut full_response = String::new();
405        let mut tool_calls: Vec<LlmToolCall> = Vec::new();
406        let max_iterations = self
407            .config
408            .providers
409            .get(&self.config.provider)
410            .map(|p| p.max_iterations)
411            .unwrap_or(100); // Allow enough iterations for complex tasks
412        let mut iteration = 0;
413        let mut consecutive_no_exec = 0;
414        let mut total_input_tokens: u64 = 0;
415        let mut total_output_tokens: u64 = 0;
416
417        while max_iterations == 0 || iteration < max_iterations {
418            iteration += 1;
419            debug!("Agent loop iteration {}", iteration);
420
421            // Send thinking event
422            debug!(
423                "Sending Thinking event with operation_id={}",
424                self.operation_id
425            );
426            self.send_event(AgentEvent::Thinking {
427                operation_id: self.operation_id,
428            });
429
430            let request_start = std::time::Instant::now();
431
432            self.maybe_compact(_messages).await;
433
434            let cached_messages = apply_cache_control(_messages, &self.config.cache);
435            let cache_count = cached_messages
436                .iter()
437                .filter(|m| m.cache_control.is_some())
438                .count();
439            debug!(
440                "Cache control applied to {} of {} messages",
441                cache_count,
442                cached_messages.len()
443            );
444
445            let mut stream = self
446                .llm_client
447                .send(cached_messages, tool_definitions.clone())
448                .await
449                .map_err(|e| CliError::ConfigError(e.to_string()))?;
450
451            tool_calls.clear();
452            let mut current_content = String::new();
453            // Track tool calls: (id) -> (name, args)
454            let mut accumulated_calls: std::collections::HashMap<
455                String,
456                (String, serde_json::Value),
457            > = std::collections::HashMap::new();
458
459            self.send_event(AgentEvent::ResponseStart {
460                operation_id: self.operation_id,
461            });
462
463            // Process stream chunks with cancellation support
464            loop {
465                // Check for cancellation FIRST (before waiting for stream)
466                if let Some(ref token) = self.cancellation_token {
467                    if token.is_cancelled() {
468                        debug!("Operation cancelled by user (pre-stream check)");
469                        self.send_event(AgentEvent::Cancelled {
470                            operation_id: self.operation_id,
471                        });
472                        return Err(CliError::ConfigError(
473                            "Operation cancelled by user".to_string(),
474                        ));
475                    }
476                }
477
478                // Use tokio::select! to check cancellation while waiting for stream
479                // Using cancellation_token.cancelled() for immediate cancellation detection
480                let chunk_result = if let Some(ref token) = self.cancellation_token {
481                    tokio::select! {
482                        chunk = stream.next() => chunk,
483                        _ = token.cancelled() => {
484                            debug!("Operation cancelled via token while waiting for stream");
485                            self.send_event(AgentEvent::Cancelled {
486                                operation_id: self.operation_id,
487                            });
488                            return Err(CliError::ConfigError("Operation cancelled by user".to_string()));
489                        }
490                    }
491                } else {
492                    stream.next().await
493                };
494
495                let Some(chunk_result) = chunk_result else {
496                    // Stream ended
497                    break;
498                };
499
500                match chunk_result {
501                    Ok(ProviderResponseChunk::ContentDelta(text)) => {
502                        current_content.push_str(&text);
503                        trace!(
504                            "ContentDelta: {} chars (total: {})",
505                            text.len(),
506                            current_content.len()
507                        );
508                        self.send_event(AgentEvent::ContentChunk {
509                            operation_id: self.operation_id,
510                            chunk: text,
511                        });
512                    }
513                    Ok(ProviderResponseChunk::ReasoningDelta(_)) => {
514                        // Ignore reasoning chunks for now
515                    }
516                    Ok(ProviderResponseChunk::ToolCallDelta {
517                        id,
518                        name,
519                        arguments,
520                    }) => {
521                        trace!(
522                            "ToolCallDelta: id={}, name={}, args_len={}",
523                            id,
524                            name,
525                            arguments.to_string().len()
526                        );
527                        // Store/merge tool call arguments
528                        accumulated_calls.insert(id.clone(), (name.clone(), arguments.clone()));
529                    }
530                    Ok(ProviderResponseChunk::Done(usage)) => {
531                        let duration_ms = request_start.elapsed().as_millis() as u64;
532                        let cost =
533                            calculate_cost(self.model(), usage.input_tokens, usage.output_tokens);
534
535                        if usage.cache_read_tokens > 0 || usage.cache_write_tokens > 0 {
536                            debug!(
537                                "Cache tokens: read={}, write={}, input={}, output={}",
538                                usage.cache_read_tokens,
539                                usage.cache_write_tokens,
540                                usage.input_tokens,
541                                usage.output_tokens
542                            );
543                        } else {
544                            debug!(
545                                "No cache tokens in response: input={}, output={}",
546                                usage.input_tokens, usage.output_tokens
547                            );
548                        }
549
550                        let _ = self.tracking_db.track_request(
551                            self.model(),
552                            usage.input_tokens,
553                            usage.output_tokens,
554                            usage.cache_read_tokens,
555                            usage.cache_write_tokens,
556                            cost,
557                            duration_ms,
558                        );
559                        total_input_tokens += usage.input_tokens;
560                        total_output_tokens += usage.output_tokens;
561                        // Emit token usage event for TUI display
562                        self.send_event(AgentEvent::TokenUsage {
563                            operation_id: self.operation_id,
564                            input_tokens: usage.input_tokens,
565                            output_tokens: usage.output_tokens,
566                        });
567                        break;
568                    }
569                    Err(e) => {
570                        let error_msg = format!("LLM error: {}", e);
571                        self.send_event(AgentEvent::Error {
572                            operation_id: self.operation_id,
573                            message: error_msg.clone(),
574                        });
575                        return Err(CliError::ConfigError(error_msg));
576                    }
577                }
578            }
579
580            // Convert accumulated calls to Vec<ToolCall> and filter invalid ones
581            let raw_tool_calls: Vec<LlmToolCall> = accumulated_calls
582                .into_iter()
583                .map(|(id, (name, args))| LlmToolCall {
584                    id,
585                    tool_type: "function".to_string(),
586                    function: limit_llm::types::FunctionCall {
587                        name,
588                        arguments: args.to_string(),
589                    },
590                })
591                .collect();
592
593            // Filter out invalid tool calls (empty names or unregistered tools)
594            let raw_count = raw_tool_calls.len();
595            tool_calls = raw_tool_calls
596                .into_iter()
597                .filter(|tc| {
598                    let is_valid = !tc.function.name.is_empty()
599                        && self.tool_names.contains(&tc.function.name.as_str());
600                    if !is_valid {
601                        debug!(
602                            "Filtered invalid tool call: id={}, name='{}'",
603                            tc.id, tc.function.name
604                        );
605                    }
606                    is_valid
607                })
608                .collect();
609
610            if tool_calls.len() != raw_count {
611                debug!(
612                    "Filtered {}/{} tool calls (empty names or unregistered tools)",
613                    raw_count - tool_calls.len(),
614                    raw_count
615                );
616            }
617
618            // BUG FIX: Don't accumulate content across iterations
619            // Only store content from the current iteration
620            // If there are tool calls, we'll continue the loop and the LLM will see the tool results
621            // If there are NO tool calls, this is the final response
622            full_response = current_content.clone();
623
624            trace!(
625                "After iter {}: content.len()={}, tool_calls={}, response.len()={}",
626                iteration,
627                current_content.len(),
628                tool_calls.len(),
629                full_response.len()
630            );
631
632            // If no tool calls, we're done
633            if tool_calls.is_empty() {
634                debug!("No tool calls, breaking loop after iteration {}", iteration);
635                break;
636            }
637
638            trace!(
639                "Tool calls found (count={}), continuing to iteration {}",
640                tool_calls.len(),
641                iteration + 1
642            );
643
644            // Execute tool calls - add assistant message with tool_calls
645            // Note: Per OpenAI API spec, when tool_calls are present, content should be null
646            let assistant_message = Message {
647                role: Role::Assistant,
648                content: None, // Don't include content when tool_calls are present
649                tool_calls: Some(tool_calls.clone()),
650                tool_call_id: None,
651                cache_control: None,
652            };
653            _messages.push(assistant_message);
654
655            // Check for duplicate tool calls and filter them out
656            let mut filtered_calls = Vec::new();
657            let mut duplicate_calls = Vec::new();
658            let mut calls_to_record = Vec::new();
659            for tc in &tool_calls {
660                let args: serde_json::Value =
661                    serde_json::from_str(&tc.function.arguments).unwrap_or_default();
662                if self.check_duplicate_tool_call(&tc.function.name, &args) {
663                    duplicate_calls.push((tc.id.clone(), tc.function.name.clone(), args));
664                } else {
665                    calls_to_record.push((tc.function.name.clone(), args.clone()));
666                    filtered_calls.push(tc.clone());
667                }
668            }
669
670            // Record new tool calls
671            for (name, args) in calls_to_record {
672                self.record_tool_call(&name, &args);
673            }
674
675            // Report duplicate calls to the model
676            if !duplicate_calls.is_empty() {
677                for (id, name, args) in &duplicate_calls {
678                    debug!(
679                        "Duplicate tool call blocked: {} with args: {}",
680                        name,
681                        serde_json::to_string(&args).unwrap_or_default()
682                    );
683                    self.send_event(AgentEvent::ToolStart {
684                        operation_id: self.operation_id,
685                        name: name.clone(),
686                        args: args.clone(),
687                    });
688                    let duplicate_msg = json!({
689                        "error": "DUPLICATE_CALL_BLOCKED",
690                        "message": format!(
691                            "You already called {} with these exact arguments in a recent turn. \
692                            Check your conversation history for the previous result. \
693                            Do not repeat the same query - use the existing data instead.",
694                            name
695                        ),
696                        "tool": name,
697                        "args": args
698                    });
699                    let result_str = serde_json::to_string(&duplicate_msg).unwrap_or_default();
700                    self.send_event(AgentEvent::ToolComplete {
701                        operation_id: self.operation_id,
702                        name: name.clone(),
703                        result: result_str.clone(),
704                    });
705                    let tool_result_message = Message {
706                        role: Role::Tool,
707                        content: Some(result_str),
708                        tool_calls: None,
709                        tool_call_id: Some(id.clone()),
710                        cache_control: None,
711                    };
712                    _messages.push(tool_result_message);
713                }
714            }
715
716            // Send ToolStart event for each non-duplicate tool BEFORE execution
717            for tc in &filtered_calls {
718                let args: serde_json::Value =
719                    serde_json::from_str(&tc.function.arguments).unwrap_or_default();
720                debug!(
721                    "ToolStart: {} with args: {}",
722                    tc.function.name,
723                    serde_json::to_string(&args).unwrap_or_default()
724                );
725                self.send_event(AgentEvent::ToolStart {
726                    operation_id: self.operation_id,
727                    name: tc.function.name.clone(),
728                    args,
729                });
730            }
731            // Execute tools (only non-duplicates)
732            let filtered_executor_calls: Vec<ToolCall> = filtered_calls
733                .iter()
734                .map(|tc| {
735                    let args: serde_json::Value =
736                        serde_json::from_str(&tc.function.arguments).unwrap_or_default();
737                    ToolCall::new(&tc.id, &tc.function.name, args)
738                })
739                .collect();
740            let results = self.executor.execute_tools(filtered_executor_calls).await;
741            let results_count = results.len();
742
743            // Add tool results to messages (OpenAI format: role=tool, tool_call_id, content)
744            for result in results {
745                let tool_call = filtered_calls.iter().find(|tc| tc.id == result.call_id);
746                if let Some(tool_call) = tool_call {
747                    let output_json = match &result.output {
748                        Ok(value) => {
749                            serde_json::to_string(value).unwrap_or_else(|_| "{}".to_string())
750                        }
751                        Err(e) => json!({ "error": e.to_string() }).to_string(),
752                    };
753
754                    debug!(
755                        "ToolComplete: {} result ({} chars): {}",
756                        tool_call.function.name,
757                        output_json.len(),
758                        output_json
759                    );
760
761                    self.send_event(AgentEvent::ToolComplete {
762                        operation_id: self.operation_id,
763                        name: tool_call.function.name.clone(),
764                        result: output_json.clone(),
765                    });
766
767                    // OpenAI tool result format
768                    let tool_result_message = Message {
769                        role: Role::Tool,
770                        content: Some(truncate_tool_result(&output_json)),
771                        tool_calls: None,
772                        tool_call_id: Some(result.call_id),
773                        cache_control: None,
774                    };
775                    _messages.push(tool_result_message);
776                }
777            }
778
779            // Safety valve: break if too many consecutive iterations without tool execution
780            if results_count == 0 && !tool_calls.is_empty() {
781                consecutive_no_exec += 1;
782                if consecutive_no_exec >= 3 {
783                    debug!(
784                        "Safety valve: {} consecutive iterations with tool calls but no executions",
785                        consecutive_no_exec
786                    );
787                    break;
788                }
789            } else {
790                consecutive_no_exec = 0;
791            }
792        }
793
794        // If we hit max iterations, make one final request to get a response (no tools = forced text)
795        // IMPORTANT: Only do this if max_iterations > 0 (0 means unlimited, so we never "hit" the limit)
796        if max_iterations > 0 && iteration >= max_iterations && !_messages.is_empty() {
797            debug!("Making final LLM call after hitting max iterations (forcing text response)");
798
799            // Add constraint message to force text response
800            let constraint_message = Message {
801                role: Role::User,
802                content: Some(
803                    "We've reached the iteration limit. Please provide a summary of:\n\
804                    1. What you've completed so far\n\
805                    2. What remains to be done\n\
806                    3. Recommended next steps for the user to continue"
807                        .to_string(),
808                ),
809                tool_calls: None,
810                tool_call_id: None,
811                cache_control: None,
812            };
813            _messages.push(constraint_message);
814
815            // Send with NO tools to force text response
816            let no_tools: Vec<LlmTool> = vec![];
817            let mut stream = self
818                .llm_client
819                .send(_messages.clone(), no_tools)
820                .await
821                .map_err(|e| CliError::ConfigError(e.to_string()))?;
822
823            // BUG FIX: Replace full_response instead of appending
824            full_response.clear();
825            self.send_event(AgentEvent::ResponseStart {
826                operation_id: self.operation_id,
827            });
828            loop {
829                // Check for cancellation FIRST (before waiting for stream)
830                if let Some(ref token) = self.cancellation_token {
831                    if token.is_cancelled() {
832                        debug!("Operation cancelled by user in final loop (pre-stream check)");
833                        self.send_event(AgentEvent::Cancelled {
834                            operation_id: self.operation_id,
835                        });
836                        return Err(CliError::ConfigError(
837                            "Operation cancelled by user".to_string(),
838                        ));
839                    }
840                }
841
842                // Use tokio::select! to check cancellation while waiting for stream
843                // Using cancellation_token.cancelled() for immediate cancellation detection
844                let chunk_result = if let Some(ref token) = self.cancellation_token {
845                    tokio::select! {
846                        chunk = stream.next() => chunk,
847                        _ = token.cancelled() => {
848                            debug!("Operation cancelled via token while waiting for stream");
849                            self.send_event(AgentEvent::Cancelled {
850                                operation_id: self.operation_id,
851                            });
852                            return Err(CliError::ConfigError("Operation cancelled by user".to_string()));
853                        }
854                    }
855                } else {
856                    stream.next().await
857                };
858
859                let Some(chunk_result) = chunk_result else {
860                    // Stream ended
861                    break;
862                };
863
864                match chunk_result {
865                    Ok(ProviderResponseChunk::ContentDelta(text)) => {
866                        full_response.push_str(&text);
867                        self.send_event(AgentEvent::ContentChunk {
868                            operation_id: self.operation_id,
869                            chunk: text,
870                        });
871                    }
872                    Ok(ProviderResponseChunk::Done(_)) => {
873                        break;
874                    }
875                    Err(e) => {
876                        debug!("Error in final LLM call: {}", e);
877                        break;
878                    }
879                    _ => {}
880                }
881            }
882        }
883
884        // IMPORTANT: Add final assistant response to message history for session persistence
885        // This is crucial for session export/share to work correctly
886        // Only add if we have content AND we haven't already added this response
887        if !full_response.is_empty() {
888            // Find the last assistant message and check if it has content
889            // If it has tool_calls but no content, UPDATE it instead of adding a new one
890            // This prevents accumulation of empty assistant messages in the history
891            let last_assistant_idx = _messages.iter().rposition(|m| m.role == Role::Assistant);
892
893            if let Some(idx) = last_assistant_idx {
894                let last_assistant = &mut _messages[idx];
895
896                // If the last assistant message has no content (tool_calls only), update it
897                if last_assistant.content.is_none()
898                    || last_assistant
899                        .content
900                        .as_ref()
901                        .map(|c| c.is_empty())
902                        .unwrap_or(true)
903                {
904                    last_assistant.content = Some(full_response.clone());
905                    debug!("Updated last assistant message with final response content");
906                } else {
907                    // Last assistant already has content, this shouldn't happen normally
908                    // but we add a new message to be safe
909                    debug!("Last assistant already has content, adding new message");
910                    let final_assistant_message = Message {
911                        role: Role::Assistant,
912                        content: Some(full_response.clone()),
913                        tool_calls: None,
914                        tool_call_id: None,
915                        cache_control: None,
916                    };
917                    _messages.push(final_assistant_message);
918                }
919            } else {
920                // No assistant message found, add a new one
921                debug!("No assistant message found, adding new message");
922                let final_assistant_message = Message {
923                    role: Role::Assistant,
924                    content: Some(full_response.clone()),
925                    tool_calls: None,
926                    tool_call_id: None,
927                    cache_control: None,
928                };
929                _messages.push(final_assistant_message);
930            }
931        }
932
933        self.send_event(AgentEvent::Done {
934            operation_id: self.operation_id,
935        });
936        Ok(ProcessResult {
937            response: full_response,
938            input_tokens: total_input_tokens,
939            output_tokens: total_output_tokens,
940        })
941    }
942
943    /// Get tool definitions formatted for the LLM
944    pub fn get_tool_definitions(&self) -> Vec<LlmTool> {
945        self.tool_names
946            .iter()
947            .map(|name| {
948                let (description, parameters) = Self::get_tool_schema(name);
949                LlmTool {
950                    tool_type: "function".to_string(),
951                    function: limit_llm::types::ToolFunction {
952                        name: name.to_string(),
953                        description,
954                        parameters,
955                    },
956                }
957            })
958            .collect()
959    }
960
961    /// Get the schema (description and parameters) for a tool
962    fn get_tool_schema(name: &str) -> (String, serde_json::Value) {
963        match name {
964            "file_read" => (
965                "Read the contents of a file".to_string(),
966                json!({
967                    "type": "object",
968                    "properties": {
969                        "path": {
970                            "type": "string",
971                            "description": "Path to the file to read"
972                        }
973                    },
974                    "required": ["path"]
975                }),
976            ),
977            "file_write" => (
978                "Write content to a file, creating parent directories if needed".to_string(),
979                json!({
980                    "type": "object",
981                    "properties": {
982                        "path": {
983                            "type": "string",
984                            "description": "Path to the file to write"
985                        },
986                        "content": {
987                            "type": "string",
988                            "description": "Content to write to the file"
989                        }
990                    },
991                    "required": ["path", "content"]
992                }),
993            ),
994            "file_edit" => (
995                "Replace text in a file with new text".to_string(),
996                json!({
997                    "type": "object",
998                    "properties": {
999                        "path": {
1000                            "type": "string",
1001                            "description": "Path to the file to edit"
1002                        },
1003                        "old_text": {
1004                            "type": "string",
1005                            "description": "Text to find and replace"
1006                        },
1007                        "new_text": {
1008                            "type": "string",
1009                            "description": "New text to replace with"
1010                        }
1011                    },
1012                    "required": ["path", "old_text", "new_text"]
1013                }),
1014            ),
1015            "bash" => (
1016                "Execute a bash command in a shell".to_string(),
1017                json!({
1018                    "type": "object",
1019                    "properties": {
1020                        "command": {
1021                            "type": "string",
1022                            "description": "Bash command to execute"
1023                        },
1024                        "workdir": {
1025                            "type": "string",
1026                            "description": "Working directory (default: current directory)"
1027                        },
1028                        "timeout": {
1029                            "type": "integer",
1030                            "description": "Timeout in seconds (default: 60)"
1031                        }
1032                    },
1033                    "required": ["command"]
1034                }),
1035            ),
1036            "git_status" => (
1037                "Get git repository status".to_string(),
1038                json!({
1039                    "type": "object",
1040                    "properties": {},
1041                    "required": []
1042                }),
1043            ),
1044            "git_diff" => (
1045                "Get git diff".to_string(),
1046                json!({
1047                    "type": "object",
1048                    "properties": {},
1049                    "required": []
1050                }),
1051            ),
1052            "git_log" => (
1053                "Get git commit log".to_string(),
1054                json!({
1055                    "type": "object",
1056                    "properties": {
1057                        "count": {
1058                            "type": "integer",
1059                            "description": "Number of commits to show (default: 10)"
1060                        }
1061                    },
1062                    "required": []
1063                }),
1064            ),
1065            "git_add" => (
1066                "Add files to git staging area".to_string(),
1067                json!({
1068                    "type": "object",
1069                    "properties": {
1070                        "files": {
1071                            "type": "array",
1072                            "items": {"type": "string"},
1073                            "description": "List of file paths to add"
1074                        }
1075                    },
1076                    "required": ["files"]
1077                }),
1078            ),
1079            "git_commit" => (
1080                "Create a git commit".to_string(),
1081                json!({
1082                    "type": "object",
1083                    "properties": {
1084                        "message": {
1085                            "type": "string",
1086                            "description": "Commit message"
1087                        }
1088                    },
1089                    "required": ["message"]
1090                }),
1091            ),
1092            "git_push" => (
1093                "Push commits to remote repository".to_string(),
1094                json!({
1095                    "type": "object",
1096                    "properties": {
1097                        "remote": {
1098                            "type": "string",
1099                            "description": "Remote name (default: origin)"
1100                        },
1101                        "branch": {
1102                            "type": "string",
1103                            "description": "Branch name (default: current branch)"
1104                        }
1105                    },
1106                    "required": []
1107                }),
1108            ),
1109            "git_pull" => (
1110                "Pull changes from remote repository".to_string(),
1111                json!({
1112                    "type": "object",
1113                    "properties": {
1114                        "remote": {
1115                            "type": "string",
1116                            "description": "Remote name (default: origin)"
1117                        },
1118                        "branch": {
1119                            "type": "string",
1120                            "description": "Branch name (default: current branch)"
1121                        }
1122                    },
1123                    "required": []
1124                }),
1125            ),
1126            "git_clone" => (
1127                "Clone a git repository".to_string(),
1128                json!({
1129                    "type": "object",
1130                    "properties": {
1131                        "url": {
1132                            "type": "string",
1133                            "description": "Repository URL to clone"
1134                        },
1135                        "directory": {
1136                            "type": "string",
1137                            "description": "Directory to clone into (optional)"
1138                        }
1139                    },
1140                    "required": ["url"]
1141                }),
1142            ),
1143            "grep" => (
1144                "Search for text patterns in files using regex".to_string(),
1145                json!({
1146                    "type": "object",
1147                    "properties": {
1148                        "pattern": {
1149                            "type": "string",
1150                            "description": "Regex pattern to search for"
1151                        },
1152                        "path": {
1153                            "type": "string",
1154                            "description": "Path to search in (default: current directory)"
1155                        }
1156                    },
1157                    "required": ["pattern"]
1158                }),
1159            ),
1160            "ast_grep" => (
1161                "AST-aware code search and transformation. Supports search, replace, and scan commands across 25+ languages. Use meta-variables: $VAR (single node), $$$VAR (multiple nodes). Search finds patterns, replace transforms code, scan runs lint rules.".to_string(),
1162                json!({
1163                    "type": "object",
1164                    "properties": {
1165                        "command": {
1166                            "type": "string",
1167                            "enum": ["search", "replace", "scan"],
1168                            "description": "Command to execute. Default: search"
1169                        },
1170                        "pattern": {
1171                            "type": "string",
1172                            "description": "AST pattern to match (e.g., 'fn $NAME() {}'). Required for search and replace."
1173                        },
1174                        "language": {
1175                            "type": "string",
1176                            "description": "Programming language. Supported: bash, c, cpp, csharp, css, elixir, go, haskell, html, java, javascript, json, kotlin, lua, nix, php, python, ruby, rust, scala, solidity, swift, typescript, tsx, yaml. Required for search and replace."
1177                        },
1178                        "path": {
1179                            "type": "string",
1180                            "description": "Path to search in (default: current directory)"
1181                        },
1182                        "rewrite": {
1183                            "type": "string",
1184                            "description": "Replacement pattern for replace command (e.g., 'logger.info($MSG)'). Required for replace."
1185                        },
1186                        "dry_run": {
1187                            "type": "boolean",
1188                            "description": "Preview replacements without modifying files (default: false). Only for replace command."
1189                        },
1190                        "globs": {
1191                            "type": "array",
1192                            "items": {"type": "string"},
1193                            "description": "Include/exclude file patterns (e.g., ['*.rs', '!*.test.rs']). Prefix with ! to exclude."
1194                        },
1195                        "context_after": {
1196                            "type": "integer",
1197                            "description": "Show N lines after each match (default: 0). Only for search."
1198                        },
1199                        "context_before": {
1200                            "type": "integer",
1201                            "description": "Show N lines before each match (default: 0). Only for search."
1202                        },
1203                        "rule": {
1204                            "type": "string",
1205                            "description": "Path to YAML rule file for scan command."
1206                        },
1207                        "inline_rules": {
1208                            "type": "string",
1209                            "description": "Inline YAML rule text for scan command."
1210                        },
1211                        "filter": {
1212                            "type": "string",
1213                            "description": "Regex to filter rules by ID for scan command."
1214                        }
1215                    },
1216                    "required": ["pattern", "language"]
1217                }),
1218            ),
1219            "lsp" => (
1220                "Perform Language Server Protocol operations (goto_definition, find_references)"
1221                    .to_string(),
1222                json!({
1223                    "type": "object",
1224                    "properties": {
1225                        "command": {
1226                            "type": "string",
1227                            "description": "LSP command: goto_definition or find_references"
1228                        },
1229                        "file_path": {
1230                            "type": "string",
1231                            "description": "Path to the file"
1232                        },
1233                        "position": {
1234                            "type": "object",
1235                            "description": "Position in the file (line, character)",
1236                            "properties": {
1237                                "line": {"type": "integer"},
1238                                "character": {"type": "integer"}
1239                            },
1240                            "required": ["line", "character"]
1241                        }
1242                    },
1243                    "required": ["command", "file_path", "position"]
1244                }),
1245            ),
1246            "web_search" => (
1247                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()),
1248                json!({
1249                    "type": "object",
1250                    "properties": {
1251                        "query": {
1252                            "type": "string",
1253                            "description": format!("Search query. Be specific for better results (e.g., 'Rust async tutorial {}' rather than 'Rust')", chrono::Local::now().year())
1254                        },
1255                        "numResults": {
1256                            "type": "integer",
1257                            "description": "Number of results to return (default: 8, max: 20)",
1258                            "default": 8
1259                        }
1260                    },
1261                    "required": ["query"]
1262                }),
1263            ),
1264            "web_fetch" => (
1265                "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(),
1266                json!({
1267                    "type": "object",
1268                    "properties": {
1269                        "url": {
1270                            "type": "string",
1271                            "description": "URL to fetch (must start with http:// or https://)"
1272                        },
1273                        "format": {
1274                            "type": "string",
1275                            "enum": ["markdown", "text", "html"],
1276                            "default": "markdown",
1277                            "description": "Output format (default: markdown)"
1278                        }
1279                    },
1280                    "required": ["url"]
1281                }),
1282            ),
1283            "browser" => (
1284                "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(),
1285                json!({
1286                    "type": "object",
1287                    "properties": {
1288                        "action": {
1289                            "type": "string",
1290                            "enum": [
1291                                // Core
1292                                "open", "close", "snapshot",
1293                                // Interaction
1294                                "click", "dblclick", "fill", "type", "press", "hover", "select",
1295                                "focus", "check", "uncheck", "scrollintoview", "drag", "upload",
1296                                // Navigation
1297                                "back", "forward", "reload",
1298                                // Query
1299                                "screenshot", "pdf", "eval", "get", "get_attr", "get_count", "get_box", "get_styles",
1300                                "find", "is", "download",
1301                                // Waiting
1302                                "wait", "wait_for_text", "wait_for_url", "wait_for_load", "wait_for_download", "wait_for_fn", "wait_for_state",
1303                                // Tabs & Dialogs
1304                                "tab_list", "tab_new", "tab_close", "tab_select", "dialog_accept", "dialog_dismiss",
1305                                // Storage & Network
1306                                "cookies", "cookies_set", "storage_get", "storage_set", "network_requests",
1307                                // Settings
1308                                "set_viewport", "set_device", "set_geo",
1309                                // State
1310                                "scroll"
1311                            ],
1312                            "description": "Browser action to perform"
1313                        },
1314                        // Core
1315                        "url": {
1316                            "type": "string",
1317                            "description": "URL to open (required for 'open' action)"
1318                        },
1319                        // Interaction
1320                        "selector": {
1321                            "type": "string",
1322                            "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)"
1323                        },
1324                        "text": {
1325                            "type": "string",
1326                            "description": "Text to input (for fill, type actions)"
1327                        },
1328                        "key": {
1329                            "type": "string",
1330                            "description": "Key to press (required for 'press' action)"
1331                        },
1332                        "value": {
1333                            "type": "string",
1334                            "description": "Value (for select, cookies_set, storage_set)"
1335                        },
1336                        "target": {
1337                            "type": "string",
1338                            "description": "Target selector (for drag action)"
1339                        },
1340                        "files": {
1341                            "type": "array",
1342                            "items": {"type": "string"},
1343                            "description": "File paths to upload (for upload action)"
1344                        },
1345                        // Query
1346                        "path": {
1347                            "type": "string",
1348                            "description": "File path (for screenshot, pdf, download actions)"
1349                        },
1350                        "script": {
1351                            "type": "string",
1352                            "description": "JavaScript to evaluate (required for 'eval' and 'wait_for_fn' actions)"
1353                        },
1354                        "get_what": {
1355                            "type": "string",
1356                            "enum": ["text", "html", "value", "url", "title"],
1357                            "description": "What to get (required for 'get' action)"
1358                        },
1359                        "attr": {
1360                            "type": "string",
1361                            "description": "Attribute name (for get_attr action)"
1362                        },
1363                        // Find
1364                        "locator_type": {
1365                            "type": "string",
1366                            "enum": ["role", "text", "label", "placeholder", "alt", "title", "testid", "css", "xpath"],
1367                            "description": "Locator strategy (for find action)"
1368                        },
1369                        "locator_value": {
1370                            "type": "string",
1371                            "description": "Locator value (for find action)"
1372                        },
1373                        "find_action": {
1374                            "type": "string",
1375                            "enum": ["click", "fill", "text", "count", "first", "last", "nth", "hover", "focus", "check", "uncheck"],
1376                            "description": "Action to perform on found element (for find action)"
1377                        },
1378                        "action_value": {
1379                            "type": "string",
1380                            "description": "Value for find action (optional)"
1381                        },
1382                        // Waiting
1383                        "wait_for": {
1384                            "type": "string",
1385                            "description": "Wait condition (for wait action)"
1386                        },
1387                        "state": {
1388                            "type": "string",
1389                            "enum": ["visible", "hidden", "attached", "detached", "enabled", "disabled", "networkidle", "domcontentloaded", "load"],
1390                            "description": "State to wait for (for wait_for_state, wait_for_load actions)"
1391                        },
1392                        // State check
1393                        "what": {
1394                            "type": "string",
1395                            "enum": ["visible", "hidden", "enabled", "disabled", "editable"],
1396                            "description": "State to check (required for 'is' action)"
1397                        },
1398                        // Scroll
1399                        "direction": {
1400                            "type": "string",
1401                            "enum": ["up", "down", "left", "right"],
1402                            "description": "Scroll direction (for scroll action)"
1403                        },
1404                        "pixels": {
1405                            "type": "integer",
1406                            "description": "Pixels to scroll (optional for scroll action)"
1407                        },
1408                        // Tabs
1409                        "index": {
1410                            "type": "integer",
1411                            "description": "Tab index (for tab_close, tab_select actions)"
1412                        },
1413                        // Dialogs
1414                        "dialog_text": {
1415                            "type": "string",
1416                            "description": "Text for prompt dialog (for dialog_accept action)"
1417                        },
1418                        // Storage
1419                        "storage_type": {
1420                            "type": "string",
1421                            "enum": ["local", "session"],
1422                            "description": "Storage type (for storage_get, storage_set actions)"
1423                        },
1424                        "key_name": {
1425                            "type": "string",
1426                            "description": "Storage key name (for storage_get, storage_set actions)"
1427                        },
1428                        // Network
1429                        "filter": {
1430                            "type": "string",
1431                            "description": "Network request filter (optional for network_requests action)"
1432                        },
1433                        // Settings
1434                        "width": {
1435                            "type": "integer",
1436                            "description": "Viewport width (for set_viewport action)"
1437                        },
1438                        "height": {
1439                            "type": "integer",
1440                            "description": "Viewport height (for set_viewport action)"
1441                        },
1442                        "scale": {
1443                            "type": "number",
1444                            "description": "Device scale factor (optional for set_viewport action)"
1445                        },
1446                        "device_name": {
1447                            "type": "string",
1448                            "description": "Device name to emulate (for set_device action)"
1449                        },
1450                        "latitude": {
1451                            "type": "number",
1452                            "description": "Latitude (for set_geo action)"
1453                        },
1454                        "longitude": {
1455                            "type": "number",
1456                            "description": "Longitude (for set_geo action)"
1457                        },
1458                        // Cookie
1459                        "name": {
1460                            "type": "string",
1461                            "description": "Cookie name (for cookies_set action)"
1462                        },
1463                        // Engine
1464                        "engine": {
1465                            "type": "string",
1466                            "enum": ["chrome", "lightpanda"],
1467                            "default": "chrome",
1468                            "description": "Browser engine to use"
1469                        }
1470                    },
1471                    "required": ["action"]
1472                }),
1473            ),
1474            _ => (
1475                format!("Tool: {}", name),
1476                json!({
1477                    "type": "object",
1478                    "properties": {},
1479                    "required": []
1480                }),
1481            ),
1482        }
1483    }
1484
1485    /// Send an event through the event channel
1486    fn send_event(&self, event: AgentEvent) {
1487        if let Some(ref tx) = self.event_tx {
1488            let _ = tx.send(event);
1489        }
1490    }
1491
1492    /// Check if the bridge is ready to process messages
1493    #[allow(dead_code)]
1494    pub fn is_ready(&self) -> bool {
1495        self.config
1496            .providers
1497            .get(&self.config.provider)
1498            .map(|p| p.api_key_or_env(&self.config.provider).is_some())
1499            .unwrap_or(false)
1500    }
1501
1502    /// Get the current model name
1503    pub fn model(&self) -> &str {
1504        self.config
1505            .providers
1506            .get(&self.config.provider)
1507            .map(|p| p.model.as_str())
1508            .unwrap_or("")
1509    }
1510
1511    /// Get the max tokens setting
1512    pub fn max_tokens(&self) -> u32 {
1513        self.config
1514            .providers
1515            .get(&self.config.provider)
1516            .map(|p| p.max_tokens)
1517            .unwrap_or(4096)
1518    }
1519
1520    /// Get the timeout setting
1521    pub fn timeout(&self) -> u64 {
1522        self.config
1523            .providers
1524            .get(&self.config.provider)
1525            .map(|p| p.timeout)
1526            .unwrap_or(60)
1527    }
1528}
1529/// Calculate cost based on model pricing (per 1M tokens)
1530fn calculate_cost(model: &str, input_tokens: u64, output_tokens: u64) -> f64 {
1531    let (input_price, output_price) = match model {
1532        // Claude 3.5 Sonnet: $3/1M input, $15/1M output
1533        "claude-3-5-sonnet-20241022" | "claude-3-5-sonnet" => (3.0, 15.0),
1534        // GPT-4: $30/1M input, $60/1M output
1535        "gpt-4" => (30.0, 60.0),
1536        // GPT-4 Turbo: $10/1M input, $30/1M output
1537        "gpt-4-turbo" | "gpt-4-turbo-preview" => (10.0, 30.0),
1538        // Default: no cost tracking
1539        _ => (0.0, 0.0),
1540    };
1541    (input_tokens as f64 * input_price / 1_000_000.0)
1542        + (output_tokens as f64 * output_price / 1_000_000.0)
1543}
1544
1545fn truncate_tool_result(result: &str) -> String {
1546    if result.len() > MAX_TOOL_RESULT_CHARS {
1547        let truncated = &result[..MAX_TOOL_RESULT_CHARS];
1548        format!(
1549            "{}\n\n... [TRUNCATED: {} chars total, showing first {}]",
1550            truncated,
1551            result.len(),
1552            MAX_TOOL_RESULT_CHARS
1553        )
1554    } else {
1555        result.to_string()
1556    }
1557}
1558
1559#[cfg(test)]
1560mod tests {
1561    use super::*;
1562    use limit_llm::{BrowserConfigSection, Config as LlmConfig, ProviderConfig};
1563    use std::collections::HashMap;
1564
1565    #[tokio::test]
1566    async fn test_agent_bridge_new() {
1567        let mut providers = HashMap::new();
1568        providers.insert(
1569            "anthropic".to_string(),
1570            ProviderConfig {
1571                api_key: Some("test-key".to_string()),
1572                model: "claude-3-5-sonnet-20241022".to_string(),
1573                base_url: None,
1574                max_tokens: 4096,
1575                timeout: 60,
1576                max_iterations: 100,
1577                thinking_enabled: false,
1578                clear_thinking: true,
1579            },
1580        );
1581        let config = LlmConfig {
1582            provider: "anthropic".to_string(),
1583            providers,
1584            browser: BrowserConfigSection::default(),
1585            compaction: limit_llm::CompactionSettings::default(),
1586            cache: limit_llm::CacheSettings::default(),
1587        };
1588
1589        let bridge = AgentBridge::new(config).unwrap();
1590        assert!(bridge.is_ready());
1591    }
1592
1593    #[tokio::test]
1594    async fn test_agent_bridge_new_no_api_key() {
1595        let mut providers = HashMap::new();
1596        providers.insert(
1597            "anthropic".to_string(),
1598            ProviderConfig {
1599                api_key: None,
1600                model: "claude-3-5-sonnet-20241022".to_string(),
1601                base_url: None,
1602                max_tokens: 4096,
1603                timeout: 60,
1604                max_iterations: 100,
1605                thinking_enabled: false,
1606                clear_thinking: true,
1607            },
1608        );
1609        let config = LlmConfig {
1610            provider: "anthropic".to_string(),
1611            providers,
1612            browser: BrowserConfigSection::default(),
1613            compaction: limit_llm::CompactionSettings::default(),
1614            cache: limit_llm::CacheSettings::default(),
1615        };
1616
1617        let result = AgentBridge::new(config);
1618        assert!(result.is_err());
1619    }
1620
1621    #[tokio::test]
1622    async fn test_get_tool_definitions() {
1623        let mut providers = HashMap::new();
1624        providers.insert(
1625            "anthropic".to_string(),
1626            ProviderConfig {
1627                api_key: Some("test-key".to_string()),
1628                model: "claude-3-5-sonnet-20241022".to_string(),
1629                base_url: None,
1630                max_tokens: 4096,
1631                timeout: 60,
1632                max_iterations: 100,
1633                thinking_enabled: false,
1634                clear_thinking: true,
1635            },
1636        );
1637        let config = LlmConfig {
1638            provider: "anthropic".to_string(),
1639            providers,
1640            browser: BrowserConfigSection::default(),
1641            compaction: limit_llm::CompactionSettings::default(),
1642            cache: limit_llm::CacheSettings::default(),
1643        };
1644
1645        let bridge = AgentBridge::new(config).unwrap();
1646        let definitions = bridge.get_tool_definitions();
1647
1648        assert_eq!(definitions.len(), 16); // grep, lsp disabled
1649
1650        // Check file_read tool definition
1651        let file_read = definitions
1652            .iter()
1653            .find(|d| d.function.name == "file_read")
1654            .unwrap();
1655        assert_eq!(file_read.tool_type, "function");
1656        assert_eq!(file_read.function.name, "file_read");
1657        assert!(file_read.function.description.contains("Read"));
1658
1659        // Check bash tool definition
1660        let bash = definitions
1661            .iter()
1662            .find(|d| d.function.name == "bash")
1663            .unwrap();
1664        assert_eq!(bash.function.name, "bash");
1665        assert!(bash.function.parameters["required"]
1666            .as_array()
1667            .unwrap()
1668            .contains(&"command".into()));
1669    }
1670
1671    #[test]
1672    fn test_get_tool_schema() {
1673        let (desc, params) = AgentBridge::get_tool_schema("file_read");
1674        assert!(desc.contains("Read"));
1675        assert_eq!(params["properties"]["path"]["type"], "string");
1676        assert!(params["required"]
1677            .as_array()
1678            .unwrap()
1679            .contains(&"path".into()));
1680
1681        let (desc, params) = AgentBridge::get_tool_schema("bash");
1682        assert!(desc.contains("bash"));
1683        assert_eq!(params["properties"]["command"]["type"], "string");
1684
1685        let (desc, _params) = AgentBridge::get_tool_schema("unknown_tool");
1686        assert!(desc.contains("unknown_tool"));
1687    }
1688
1689    #[test]
1690    fn test_is_ready() {
1691        let mut providers = HashMap::new();
1692        providers.insert(
1693            "anthropic".to_string(),
1694            ProviderConfig {
1695                api_key: Some("test-key".to_string()),
1696                model: "claude-3-5-sonnet-20241022".to_string(),
1697                base_url: None,
1698                max_tokens: 4096,
1699                timeout: 60,
1700                max_iterations: 100,
1701                thinking_enabled: false,
1702                clear_thinking: true,
1703            },
1704        );
1705        let config_with_key = LlmConfig {
1706            provider: "anthropic".to_string(),
1707            providers,
1708            browser: BrowserConfigSection::default(),
1709            compaction: limit_llm::CompactionSettings::default(),
1710            cache: limit_llm::CacheSettings::default(),
1711        };
1712
1713        let bridge = AgentBridge::new(config_with_key).unwrap();
1714        assert!(bridge.is_ready());
1715    }
1716
1717    #[test]
1718    fn test_handoff_compaction_preserves_system() {
1719        let handoff = ModelHandoff::new();
1720
1721        let mut messages = vec![Message {
1722            role: Role::System,
1723            content: Some("System prompt".to_string()),
1724            tool_calls: None,
1725            tool_call_id: None,
1726            cache_control: None,
1727        }];
1728
1729        for i in 0..50 {
1730            messages.push(Message {
1731                role: if i % 2 == 0 {
1732                    Role::User
1733                } else {
1734                    Role::Assistant
1735                },
1736                content: Some(format!(
1737                    "Message {} with enough content to consume tokens",
1738                    i
1739                )),
1740                tool_calls: None,
1741                tool_call_id: None,
1742                cache_control: None,
1743            });
1744        }
1745
1746        let target = 500;
1747        let compacted = handoff.compact_messages(&messages, target);
1748
1749        assert_eq!(compacted[0].role, Role::System);
1750        assert!(compacted.len() < messages.len());
1751    }
1752
1753    #[test]
1754    fn test_handoff_compaction_keeps_recent() {
1755        let handoff = ModelHandoff::new();
1756
1757        let mut messages = vec![Message {
1758            role: Role::System,
1759            content: Some("System".to_string()),
1760            tool_calls: None,
1761            tool_call_id: None,
1762            cache_control: None,
1763        }];
1764
1765        for i in 0..100 {
1766            messages.push(Message {
1767                role: if i % 2 == 0 {
1768                    Role::User
1769                } else {
1770                    Role::Assistant
1771                },
1772                content: Some(format!("Message {}", i)),
1773                tool_calls: None,
1774                tool_call_id: None,
1775                cache_control: None,
1776            });
1777        }
1778
1779        let target = 200;
1780        let compacted = handoff.compact_messages(&messages, target);
1781
1782        assert!(compacted.len() < messages.len());
1783        let last_content = compacted.last().unwrap().content.clone();
1784        assert_eq!(last_content, Some("Message 99".to_string()));
1785    }
1786
1787    #[test]
1788    fn test_compaction_config_respects_settings() {
1789        let mut providers = HashMap::new();
1790        providers.insert(
1791            "anthropic".to_string(),
1792            ProviderConfig {
1793                api_key: Some("test-key".to_string()),
1794                model: "claude-3-5-sonnet-20241022".to_string(),
1795                base_url: None,
1796                max_tokens: 4096,
1797                timeout: 60,
1798                max_iterations: 100,
1799                thinking_enabled: false,
1800                clear_thinking: true,
1801            },
1802        );
1803
1804        let config = LlmConfig {
1805            provider: "anthropic".to_string(),
1806            providers,
1807            browser: BrowserConfigSection::default(),
1808            compaction: limit_llm::CompactionSettings {
1809                enabled: true,
1810                reserve_tokens: 8192,
1811                keep_recent_tokens: 10000,
1812                use_summarization: true,
1813            },
1814            cache: limit_llm::CacheSettings::default(),
1815        };
1816
1817        let bridge = AgentBridge::new(config).unwrap();
1818        assert!(bridge.config.compaction.enabled);
1819        assert_eq!(bridge.config.compaction.reserve_tokens, 8192);
1820        assert_eq!(bridge.config.compaction.keep_recent_tokens, 10000);
1821    }
1822}