limit_cli/tui/bridge/
bridge_impl.rs

1//! TUI Bridge module
2//!
3//! Bridge connecting limit-cli REPL to limit-tui components
4
5use crate::agent_bridge::{AgentBridge, AgentEvent};
6use crate::error::CliError;
7use crate::session::SessionManager;
8use crate::tui::{activity::format_activity_message, TuiState};
9use limit_tui::components::{ActivityFeed, ChatView, Message, Spinner};
10use std::sync::{Arc, Mutex};
11use tokio::sync::mpsc;
12use tracing::trace;
13
14/// Bridge connecting limit-cli REPL to limit-tui components
15pub struct TuiBridge {
16    /// Agent bridge for processing messages (wrapped for thread-safe access)
17    agent_bridge: Arc<Mutex<AgentBridge>>,
18    /// Event receiver from the agent
19    event_rx: mpsc::UnboundedReceiver<AgentEvent>,
20    /// Current TUI state
21    state: Arc<Mutex<TuiState>>,
22    /// Chat view for displaying conversation
23    chat_view: Arc<Mutex<ChatView>>,
24    /// Activity feed for showing tool activities
25    activity_feed: Arc<Mutex<ActivityFeed>>,
26    /// Spinner for thinking state
27    /// Spinner for thinking state
28    spinner: Arc<Mutex<Spinner>>,
29    /// Conversation history
30    messages: Arc<Mutex<Vec<limit_llm::Message>>>,
31    /// Total input tokens for the session
32    total_input_tokens: Arc<Mutex<u64>>,
33    /// Total output tokens for the session
34    total_output_tokens: Arc<Mutex<u64>>,
35    /// Session manager for persistence
36    session_manager: Arc<Mutex<SessionManager>>,
37    /// Current session ID
38    session_id: Arc<Mutex<String>>,
39    /// Current operation ID (to ignore events from old operations)
40    operation_id: Arc<Mutex<u64>>,
41}
42
43impl TuiBridge {
44    /// Create a new TuiBridge with the given agent bridge and event channel
45    pub fn new(
46        agent_bridge: AgentBridge,
47        event_rx: mpsc::UnboundedReceiver<AgentEvent>,
48    ) -> Result<Self, CliError> {
49        let session_manager = SessionManager::new().map_err(|e| {
50            CliError::ConfigError(format!("Failed to create session manager: {}", e))
51        })?;
52
53        Self::with_session_manager(agent_bridge, event_rx, session_manager)
54    }
55
56    /// Create a new TuiBridge for testing with a temporary session manager
57    #[cfg(test)]
58    pub fn new_for_test(
59        agent_bridge: AgentBridge,
60        event_rx: mpsc::UnboundedReceiver<AgentEvent>,
61    ) -> Result<Self, CliError> {
62        use tempfile::TempDir;
63
64        // Create a temporary directory for the test
65        let temp_dir = TempDir::new().map_err(|e| {
66            CliError::ConfigError(format!("Failed to create temp directory: {}", e))
67        })?;
68
69        let db_path = temp_dir.path().join("session.db");
70        let sessions_dir = temp_dir.path().join("sessions");
71
72        let session_manager = SessionManager::with_paths(db_path, sessions_dir).map_err(|e| {
73            CliError::ConfigError(format!("Failed to create session manager: {}", e))
74        })?;
75
76        Self::with_session_manager(agent_bridge, event_rx, session_manager)
77    }
78
79    /// Create a new TuiBridge with a custom session manager
80    pub fn with_session_manager(
81        agent_bridge: AgentBridge,
82        event_rx: mpsc::UnboundedReceiver<AgentEvent>,
83        session_manager: SessionManager,
84    ) -> Result<Self, CliError> {
85        // Always create a new session on TUI startup
86        let session_id = session_manager
87            .create_new_session()
88            .map_err(|e| CliError::ConfigError(format!("Failed to create session: {}", e)))?;
89        tracing::info!("Created new TUI session: {}", session_id);
90
91        // Start with empty messages - never load previous session
92        let messages: Vec<limit_llm::Message> = Vec::new();
93
94        // Get token counts from session info
95        let sessions = session_manager.list_sessions().unwrap_or_default();
96        let session_info = sessions.iter().find(|s| s.id == session_id);
97        let initial_input = session_info.map(|s| s.total_input_tokens).unwrap_or(0);
98        let initial_output = session_info.map(|s| s.total_output_tokens).unwrap_or(0);
99
100        let chat_view = Arc::new(Mutex::new(ChatView::new()));
101
102        // Add loaded messages to chat view for display
103        for msg in &messages {
104            match msg.role {
105                limit_llm::Role::User => {
106                    let text = msg
107                        .content
108                        .as_ref()
109                        .map(|c| c.to_text())
110                        .unwrap_or_default();
111                    let chat_msg = Message::user(text);
112                    chat_view.lock().unwrap().add_message(chat_msg);
113                }
114                limit_llm::Role::Assistant => {
115                    let text = msg
116                        .content
117                        .as_ref()
118                        .map(|c| c.to_text())
119                        .unwrap_or_default();
120                    let chat_msg = Message::assistant(text);
121                    chat_view.lock().unwrap().add_message(chat_msg);
122                }
123                limit_llm::Role::System => {
124                    // Skip system messages in display
125                }
126                limit_llm::Role::Tool => {
127                    // Skip tool messages in display
128                }
129            }
130        }
131
132        tracing::info!("Loaded {} messages into chat view", messages.len());
133
134        // Add system message to indicate this is a new session
135        let session_short_id = format!("...{}", &session_id[session_id.len().saturating_sub(8)..]);
136        let welcome_msg =
137            Message::system(format!("🆕 New TUI session started: {}", session_short_id));
138        chat_view.lock().unwrap().add_message(welcome_msg);
139
140        // Add model info as system message
141        let model_name = agent_bridge.model().to_string();
142        if !model_name.is_empty() {
143            let model_msg = Message::system(format!("Using model: {}", model_name));
144            chat_view.lock().unwrap().add_message(model_msg);
145        }
146
147        Ok(Self {
148            agent_bridge: Arc::new(Mutex::new(agent_bridge)),
149            event_rx,
150            state: Arc::new(Mutex::new(TuiState::Idle)),
151            chat_view,
152            activity_feed: Arc::new(Mutex::new(ActivityFeed::new())),
153            spinner: Arc::new(Mutex::new(Spinner::new("Thinking..."))),
154            messages: Arc::new(Mutex::new(messages)),
155            total_input_tokens: Arc::new(Mutex::new(initial_input)),
156            total_output_tokens: Arc::new(Mutex::new(initial_output)),
157            session_manager: Arc::new(Mutex::new(session_manager)),
158            session_id: Arc::new(Mutex::new(session_id)),
159            operation_id: Arc::new(Mutex::new(0)),
160        })
161    }
162
163    /// Get a clone of the agent bridge Arc for spawning tasks
164    pub fn agent_bridge_arc(&self) -> Arc<Mutex<AgentBridge>> {
165        self.agent_bridge.clone()
166    }
167
168    /// Get locked access to the agent bridge (for compatibility)
169    #[allow(dead_code)]
170    pub fn agent_bridge(&self) -> std::sync::MutexGuard<'_, AgentBridge> {
171        self.agent_bridge.lock().unwrap()
172    }
173
174    /// Get the current TUI state
175    pub fn state(&self) -> TuiState {
176        self.state.lock().unwrap().clone()
177    }
178
179    /// Get a reference to the chat view
180    pub fn chat_view(&self) -> &Arc<Mutex<ChatView>> {
181        &self.chat_view
182    }
183
184    /// Get a reference to the spinner
185    pub fn spinner(&self) -> &Arc<Mutex<Spinner>> {
186        &self.spinner
187    }
188
189    /// Get a reference to the activity feed
190    pub fn activity_feed(&self) -> &Arc<Mutex<ActivityFeed>> {
191        &self.activity_feed
192    }
193
194    /// Process events from the agent and update TUI state
195    pub fn process_events(&mut self) -> Result<(), CliError> {
196        let mut event_count = 0;
197        let current_op_id = self.operation_id();
198
199        while let Ok(event) = self.event_rx.try_recv() {
200            event_count += 1;
201
202            // Get operation_id from event
203            let event_op_id = match &event {
204                AgentEvent::Thinking { operation_id } => *operation_id,
205                AgentEvent::ToolStart { operation_id, .. } => *operation_id,
206                AgentEvent::ToolComplete { operation_id, .. } => *operation_id,
207                AgentEvent::ResponseStart { operation_id } => *operation_id,
208                AgentEvent::ContentChunk { operation_id, .. } => *operation_id,
209                AgentEvent::Done { operation_id } => *operation_id,
210                AgentEvent::Cancelled { operation_id } => *operation_id,
211                AgentEvent::Error { operation_id, .. } => *operation_id,
212                AgentEvent::TokenUsage { operation_id, .. } => *operation_id,
213            };
214
215            trace!(
216                "process_events: event_op_id={}, current_op_id={}, event={:?}",
217                event_op_id,
218                current_op_id,
219                std::mem::discriminant(&event)
220            );
221
222            // Ignore events from old operations
223            if event_op_id != current_op_id {
224                trace!(
225                    "process_events: Ignoring event from old operation {} (current: {})",
226                    event_op_id,
227                    current_op_id
228                );
229                continue;
230            }
231
232            match event {
233                AgentEvent::Thinking { operation_id: _ } => {
234                    trace!("process_events: Thinking event received - setting state to Thinking",);
235                    *self.state.lock().unwrap() = TuiState::Thinking;
236                    trace!("process_events: state is now {:?}", self.state());
237                }
238                AgentEvent::ToolStart {
239                    operation_id: _,
240                    name,
241                    args,
242                } => {
243                    trace!("process_events: ToolStart event - {}", name);
244                    let activity_msg = format_activity_message(&name, &args);
245                    // Add to activity feed instead of changing state
246                    self.activity_feed.lock().unwrap().add(activity_msg, true);
247                }
248                AgentEvent::ToolComplete {
249                    operation_id: _,
250                    name: _,
251                    result: _,
252                } => {
253                    trace!("process_events: ToolComplete event");
254                    // Mark current activity as complete
255                    self.activity_feed.lock().unwrap().complete_current();
256                }
257                AgentEvent::ResponseStart { operation_id: _ } => {
258                    trace!("process_events: ResponseStart event - creating new assistant message");
259                    self.chat_view.lock().unwrap().start_new_assistant_message();
260                }
261                AgentEvent::ContentChunk {
262                    operation_id: _,
263                    chunk,
264                } => {
265                    trace!("process_events: ContentChunk event ({} chars)", chunk.len());
266                    self.chat_view
267                        .lock()
268                        .unwrap()
269                        .append_to_last_assistant(&chunk);
270                }
271                AgentEvent::Done { operation_id: _ } => {
272                    trace!("process_events: Done event received");
273                    *self.state.lock().unwrap() = TuiState::Idle;
274                    // Mark all activities as complete when LLM finishes
275                    self.activity_feed.lock().unwrap().complete_all();
276                }
277                AgentEvent::Cancelled { operation_id: _ } => {
278                    trace!("process_events: Cancelled event received");
279                    *self.state.lock().unwrap() = TuiState::Idle;
280                    // Mark all activities as complete
281                    self.activity_feed.lock().unwrap().complete_all();
282                }
283                AgentEvent::Error {
284                    operation_id: _,
285                    message,
286                } => {
287                    trace!("process_events: Error event - {}", message);
288                    // Reset state to Idle so user can continue
289                    *self.state.lock().unwrap() = TuiState::Idle;
290                    let chat_msg = Message::system(format!("Error: {}", message));
291                    self.chat_view.lock().unwrap().add_message(chat_msg);
292                }
293                AgentEvent::TokenUsage { .. } => {}
294            }
295        }
296        if event_count > 0 {
297            trace!("process_events: processed {} events", event_count);
298        }
299        Ok(())
300    }
301
302    /// Add a user message to the chat
303    pub fn add_user_message(&self, content: String) {
304        let msg = Message::user(content);
305        self.chat_view.lock().unwrap().add_message(msg);
306    }
307
308    /// Tick the spinner animation
309    pub fn tick_spinner(&self) {
310        self.spinner.lock().unwrap().tick();
311    }
312
313    /// Check if agent is busy
314    pub fn is_busy(&self) -> bool {
315        !matches!(self.state(), TuiState::Idle)
316    }
317
318    /// Get current operation ID
319    #[inline]
320    pub fn operation_id(&self) -> u64 {
321        *self.operation_id.lock().unwrap_or_else(|e| e.into_inner())
322    }
323
324    /// Increment and get new operation ID
325    pub fn next_operation_id(&self) -> u64 {
326        let mut id = self.operation_id.lock().unwrap_or_else(|e| e.into_inner());
327        *id += 1;
328        *id
329    }
330
331    /// Get total input tokens for the session
332    #[inline]
333    pub fn total_input_tokens(&self) -> u64 {
334        *self
335            .total_input_tokens
336            .lock()
337            .unwrap_or_else(|e| e.into_inner())
338    }
339
340    /// Get total output tokens for the session
341    #[inline]
342    pub fn total_output_tokens(&self) -> u64 {
343        *self
344            .total_output_tokens
345            .lock()
346            .unwrap_or_else(|e| e.into_inner())
347    }
348
349    /// Get the current session ID
350    pub fn session_id(&self) -> String {
351        self.session_id
352            .lock()
353            .map(|guard| guard.clone())
354            .unwrap_or_else(|_| String::from("unknown"))
355    }
356
357    /// Save the current session
358    pub fn save_session(&self) -> Result<(), CliError> {
359        let session_id = self
360            .session_id
361            .lock()
362            .map(|guard| guard.clone())
363            .unwrap_or_else(|_| String::from("unknown"));
364
365        let messages = self
366            .messages
367            .lock()
368            .map(|guard| guard.clone())
369            .unwrap_or_default();
370
371        let input_tokens = self
372            .total_input_tokens
373            .lock()
374            .map(|guard| *guard)
375            .unwrap_or(0);
376
377        let output_tokens = self
378            .total_output_tokens
379            .lock()
380            .map(|guard| *guard)
381            .unwrap_or(0);
382
383        tracing::debug!(
384            "Saving session {} with {} messages, {} in tokens, {} out tokens",
385            session_id,
386            messages.len(),
387            input_tokens,
388            output_tokens
389        );
390
391        let session_manager = self.session_manager.lock().map_err(|e| {
392            CliError::ConfigError(format!("Failed to acquire session manager lock: {}", e))
393        })?;
394
395        session_manager.save_session(&session_id, &messages, input_tokens, output_tokens)?;
396
397        if !messages.is_empty() {
398            if let Err(e) = session_manager.migrate_to_tree(&session_id) {
399                tracing::warn!("Failed to migrate session to tree format: {}", e);
400            }
401        }
402
403        tracing::info!(
404            "✓ Session {} saved successfully ({} messages, {} in tokens, {} out tokens)",
405            session_id,
406            messages.len(),
407            input_tokens,
408            output_tokens
409        );
410        Ok(())
411    }
412
413    /// Get session manager (for command handling)
414    pub fn session_manager(&self) -> Arc<Mutex<SessionManager>> {
415        self.session_manager.clone()
416    }
417
418    /// Get messages arc (for command handling)
419    pub fn messages(&self) -> Arc<Mutex<Vec<limit_llm::Message>>> {
420        self.messages.clone()
421    }
422
423    /// Get state arc (for command handling)
424    pub fn state_arc(&self) -> Arc<Mutex<TuiState>> {
425        self.state.clone()
426    }
427
428    /// Get total input tokens arc (for command handling)
429    pub fn total_input_tokens_arc(&self) -> Arc<Mutex<u64>> {
430        self.total_input_tokens.clone()
431    }
432
433    /// Get total output tokens arc (for command handling)
434    pub fn total_output_tokens_arc(&self) -> Arc<Mutex<u64>> {
435        self.total_output_tokens.clone()
436    }
437
438    /// Get session id arc (for command handling)
439    pub fn session_id_arc(&self) -> Arc<Mutex<String>> {
440        self.session_id.clone()
441    }
442
443    /// Set state (for cancellation)
444    pub fn set_state(&self, new_state: TuiState) {
445        *self.state.lock().unwrap() = new_state;
446    }
447}
448
449#[cfg(test)]
450mod tests {
451    use super::*;
452
453    /// Create a test config for AgentBridge
454    fn create_test_config() -> limit_llm::Config {
455        use limit_llm::{BrowserConfigSection, ProviderConfig};
456        let mut providers = std::collections::HashMap::new();
457        providers.insert(
458            "anthropic".to_string(),
459            ProviderConfig {
460                api_key: Some("test-key".to_string()),
461                model: "claude-3-5-sonnet-20241022".to_string(),
462                base_url: None,
463                max_tokens: 4096,
464                timeout: 60,
465                max_iterations: 100,
466                thinking_enabled: false,
467                clear_thinking: true,
468            },
469        );
470        limit_llm::Config {
471            provider: "anthropic".to_string(),
472            providers,
473            browser: BrowserConfigSection::default(),
474            compaction: limit_llm::CompactionSettings::default(),
475            cache: limit_llm::CacheSettings::default(),
476        }
477    }
478
479    #[test]
480    fn test_tui_bridge_new() {
481        let config = create_test_config();
482        let agent_bridge = AgentBridge::new(config).unwrap();
483        let (_tx, rx) = mpsc::unbounded_channel();
484
485        let tui_bridge = TuiBridge::new(agent_bridge, rx).unwrap();
486        assert_eq!(tui_bridge.state(), TuiState::Idle);
487    }
488
489    #[test]
490    fn test_tui_bridge_state() {
491        let config = create_test_config();
492        let agent_bridge = AgentBridge::new(config).unwrap();
493        let (tx, rx) = mpsc::unbounded_channel();
494
495        let mut tui_bridge = TuiBridge::new(agent_bridge, rx).unwrap();
496
497        let op_id = tui_bridge.operation_id();
498        tx.send(AgentEvent::Thinking {
499            operation_id: op_id,
500        })
501        .unwrap();
502        tui_bridge.process_events().unwrap();
503        assert!(matches!(tui_bridge.state(), TuiState::Thinking));
504
505        tx.send(AgentEvent::Done {
506            operation_id: op_id,
507        })
508        .unwrap();
509        tui_bridge.process_events().unwrap();
510        assert_eq!(tui_bridge.state(), TuiState::Idle);
511    }
512
513    #[test]
514    fn test_tui_bridge_chat_view() {
515        let config = create_test_config();
516        let agent_bridge = AgentBridge::new(config).unwrap();
517        let (_tx, rx) = mpsc::unbounded_channel();
518
519        let tui_bridge = TuiBridge::new(agent_bridge, rx).unwrap();
520
521        tui_bridge.add_user_message("Hello".to_string());
522        assert_eq!(tui_bridge.chat_view().lock().unwrap().message_count(), 3); // 1 user + 2 system (welcome + model)
523    }
524}
limit_cli/tui/bridge/bridge_impl.rs

limit_cli/tui/bridge/
bridge_impl.rs
