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

limit_cli/
agent_bridge.rs
