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

limit_cli/tui/bridge/
bridge_impl.rs
