mermaid_cli/tui/
app.rs

1//! Application coordinator
2//!
3//! Thin coordinator that composes state modules. All state is delegated to
4//! focused modules in src/tui/state/.
5
6use std::collections::VecDeque;
7use std::sync::Arc;
8use tokio::task::JoinHandle;
9
10use super::state::{
11    AppState, AttachmentState, ConversationState, ErrorEntry, ErrorSeverity, GenerationStatus,
12    InputBuffer, ModelState, OperationState, StatusState, UIState,
13};
14use crate::constants::UI_ERROR_LOG_MAX_SIZE;
15use crate::models::{ChatMessage, MessageRole, Model, ModelConfig, StreamCallback};
16use crate::session::{ConversationHistory, ConversationManager};
17
18/// Application state coordinator
19pub struct App {
20    /// User input buffer
21    pub input: InputBuffer,
22    /// Is the app running?
23    pub running: bool,
24    /// Current working directory
25    pub working_dir: String,
26    /// Error log - keeps last N errors for visibility
27    pub error_log: VecDeque<ErrorEntry>,
28    /// State machine for application lifecycle
29    pub app_state: AppState,
30
31    /// Model state - LLM configuration
32    pub model_state: ModelState,
33    /// UI state - visual presentation and widget states
34    pub ui_state: UIState,
35    /// Session state - conversation history and persistence
36    pub session_state: ConversationState,
37    /// Operation state - file reading and tool calls
38    pub operation_state: OperationState,
39    /// Status state - UI status messages
40    pub status_state: StatusState,
41    /// Attachment state - pending image attachments
42    pub attachment_state: AttachmentState,
43    /// MCP tool definitions in Ollama JSON format (injected at startup or background init)
44    pub mcp_tools: Vec<serde_json::Value>,
45    /// Background MCP server initialization task.
46    /// Polled in the event loop; when done, mcp_tools and global manager are set.
47    pub mcp_init_task: Option<JoinHandle<McpInitResult>>,
48}
49
50/// Result of background MCP server initialization
51pub struct McpInitResult {
52    pub tools: Vec<serde_json::Value>,
53    pub manager: Option<Arc<crate::mcp::McpServerManager>>,
54}
55
56impl App {
57    /// Create a new app instance
58    pub fn new(model: Box<dyn Model>, model_id: String, base_config: ModelConfig) -> Self {
59        let working_dir = std::env::current_dir()
60            .map(|p| p.to_string_lossy().to_string())
61            .unwrap_or_else(|_| ".".to_string());
62
63        // Initialize model state
64        let model_state = ModelState::new(model, model_id, base_config);
65
66        // Initialize conversation manager for the current directory
67        let conversation_manager = ConversationManager::new(&working_dir).ok();
68        let current_conversation = conversation_manager
69            .as_ref()
70            .map(|_| ConversationHistory::new(working_dir.clone(), model_state.model_name.clone()));
71
72        // Load input history from conversation if available
73        let input_history: std::collections::VecDeque<String> = current_conversation
74            .as_ref()
75            .map(|conv| conv.input_history.clone())
76            .unwrap_or_default();
77
78        // Initialize input buffer with persisted history
79        let mut input = InputBuffer::new();
80        input.load_history(input_history);
81
82        // Initialize UIState
83        let ui_state = UIState::new();
84
85        // Initialize ConversationState with conversation management
86        let session_state = ConversationState::with_conversation(
87            conversation_manager,
88            current_conversation,
89        );
90
91        Self {
92            input,
93            running: true,
94            working_dir,
95            error_log: VecDeque::new(),
96            app_state: AppState::Idle,
97            model_state,
98            ui_state,
99            session_state,
100            operation_state: OperationState::new(),
101            status_state: StatusState::new(),
102            attachment_state: AttachmentState::new(),
103            mcp_tools: Vec::new(),
104            mcp_init_task: None,
105        }
106    }
107
108    /// Build a ModelConfig with MCP tools included.
109    pub fn build_model_config(&self) -> crate::models::ModelConfig {
110        let mut config = self.model_state.build_config();
111        config.mcp_tools = self.mcp_tools.clone();
112        config
113    }
114
115    /// Poll for completed MCP background initialization (non-blocking).
116    ///
117    /// Returns immediately if init is still in progress or already done.
118    /// When init completes, sets mcp_tools on self and registers the global
119    /// MCP manager so tool calls can be dispatched.
120    ///
121    /// Called from both the main event loop and the agent loop so that MCP
122    /// tools become available to the model as soon as servers are ready,
123    /// even mid-agent-loop.
124    pub async fn poll_mcp_init(&mut self) {
125        let ready = self.mcp_init_task.as_ref().is_some_and(|t| t.is_finished());
126        if !ready {
127            return;
128        }
129        if let Some(task) = self.mcp_init_task.take()
130            && let Ok(result) = task.await
131        {
132            if !result.tools.is_empty() {
133                self.mcp_tools = result.tools;
134            }
135            if let Some(manager) = result.manager {
136                crate::agents::set_mcp_manager(manager);
137            }
138        }
139        // Mark complete whether the task succeeded or panicked, so waiters unblock
140        crate::agents::mark_mcp_init_complete();
141    }
142
143    // ===== Message Management =====
144
145    /// Add a message to the chat (extracts thinking blocks automatically)
146    pub fn add_message(&mut self, role: MessageRole, content: String) {
147        self.add_message_with_images(role, content, None);
148    }
149
150    /// Add a message with optional image attachments
151    pub fn add_message_with_images(
152        &mut self,
153        role: MessageRole,
154        content: String,
155        images: Option<Vec<String>>,
156    ) {
157        let mut message = match role {
158            MessageRole::User => ChatMessage::user(content),
159            MessageRole::Assistant => ChatMessage::assistant(content),
160            MessageRole::System => ChatMessage::system(content),
161            MessageRole::Tool => ChatMessage::tool("", "", content),
162        };
163        let (thinking, answer) = ChatMessage::extract_thinking(&message.content);
164        message.content = answer;
165        message.thinking = thinking;
166        if let Some(imgs) = images {
167            message = message.with_images(imgs);
168        }
169        self.commit_message(message);
170    }
171
172    /// Add an assistant message with tool_calls attached
173    pub fn add_assistant_message_with_tool_calls(
174        &mut self,
175        content: String,
176        tool_calls: Vec<crate::models::ToolCall>,
177    ) {
178        let mut message = ChatMessage::assistant(content).with_tool_calls(tool_calls);
179        let (thinking, answer) = ChatMessage::extract_thinking(&message.content);
180        message.content = answer;
181        message.thinking = thinking;
182        self.commit_message(message);
183    }
184
185    /// Add a tool result message
186    pub fn add_tool_result(&mut self, tool_call_id: String, tool_name: String, content: String) {
187        let message = ChatMessage::tool(tool_call_id, tool_name, content);
188        self.commit_message(message);
189    }
190
191    /// Commit a message to session state and conversation history
192    pub fn commit_message(&mut self, message: ChatMessage) {
193        self.session_state.messages.push(message.clone());
194        if let Some(ref mut conv) = self.session_state.current_conversation {
195            conv.add_messages(&[message]);
196        }
197    }
198
199    /// Clear the input buffer
200    pub fn clear_input(&mut self) {
201        self.input.clear();
202    }
203
204    // ===== Status Management =====
205
206    /// Set status message
207    pub fn set_status(&mut self, message: impl Into<String>) {
208        self.status_state.set(message);
209    }
210
211    /// Clear status message
212    pub fn clear_status(&mut self) {
213        self.status_state.clear();
214    }
215
216    // ===== Error Management =====
217
218    /// Display an error consistently across the UI
219    pub fn display_error(&mut self, summary: impl Into<String>, detail: impl Into<String>) {
220        let summary = summary.into();
221        let detail = detail.into();
222
223        self.set_status(format!("[Error] {}", summary));
224
225        if detail.is_empty() {
226            self.add_message(MessageRole::System, format!("Error: {}", summary));
227        } else {
228            self.add_message(MessageRole::System, detail);
229        }
230    }
231
232    /// Display an error with just a message
233    pub fn display_error_simple(&mut self, message: impl Into<String>) {
234        let message = message.into();
235        self.display_error(message.clone(), message);
236    }
237
238    /// Log an error to the error log
239    pub fn log_error(&mut self, entry: ErrorEntry) {
240        self.status_state.set(entry.display());
241        self.error_log.push_back(entry);
242        if self.error_log.len() > UI_ERROR_LOG_MAX_SIZE {
243            self.error_log.pop_front(); // O(1) instead of O(n)
244        }
245    }
246
247    /// Log a simple error message
248    pub fn log_error_msg(&mut self, severity: ErrorSeverity, msg: impl Into<String>) {
249        self.log_error(ErrorEntry::new(severity, msg.into()));
250    }
251
252    /// Log error with context
253    pub fn log_error_with_context(
254        &mut self,
255        severity: ErrorSeverity,
256        msg: impl Into<String>,
257        context: impl Into<String>,
258    ) {
259        self.log_error(ErrorEntry::with_context(
260            severity,
261            msg.into(),
262            context.into(),
263        ));
264    }
265
266    /// Get recent errors
267    pub fn recent_errors(&self, count: usize) -> Vec<&ErrorEntry> {
268        self.error_log.iter().rev().take(count).collect()
269    }
270
271    // ===== Terminal =====
272
273    /// Set terminal window title
274    pub fn set_terminal_title(&self, title: &str) {
275        use crossterm::{execute, terminal::SetTitle};
276        use std::io::stdout;
277        let _ = execute!(stdout(), SetTitle(title));
278    }
279
280    // ===== Title Generation =====
281
282    /// Spawn title generation as a background task (non-blocking).
283    /// Returns a JoinHandle the caller can poll with `is_finished()`.
284    pub fn spawn_title_generation(&self) -> Option<tokio::task::JoinHandle<Option<String>>> {
285        if self.session_state.conversation_title.is_some() || self.session_state.messages.len() < 2
286        {
287            return None;
288        }
289
290        let mut summary = String::new();
291        for msg in self
292            .session_state
293            .messages
294            .iter()
295            .filter(|m| matches!(m.role, MessageRole::User | MessageRole::Assistant))
296            .take(4)
297        {
298            let role = if msg.role == MessageRole::User {
299                "User"
300            } else {
301                "Assistant"
302            };
303            summary.push_str(&format!(
304                "{}: {}\n\n",
305                role,
306                msg.content.chars().take(200).collect::<String>()
307            ));
308        }
309
310        let model = self.model_state.model.clone();
311        let mut config = self.build_model_config();
312        config.thinking_enabled = Some(false);
313
314        Some(tokio::spawn(async move {
315            let prompt = format!(
316                "Based on this conversation, generate a short, descriptive title (2-4 words maximum, no quotes):\n\n{}\n\nTitle:",
317                summary
318            );
319            let buf = Arc::new(tokio::sync::Mutex::new(String::new()));
320            let buf_clone = Arc::clone(&buf);
321            let callback: StreamCallback = Arc::new(move |chunk: &str| {
322                if let Ok(mut t) = buf_clone.try_lock() {
323                    t.push_str(chunk);
324                }
325            });
326
327            let model = model.read().await;
328            if model
329                .chat(&[ChatMessage::user(prompt)], &config, Some(callback))
330                .await
331                .is_ok()
332            {
333                let raw = buf.lock().await;
334                let title: String = raw
335                    .lines()
336                    .next()
337                    .unwrap_or(&raw)
338                    .trim()
339                    .trim_matches(|c| c == '"' || c == '\'' || c == '.' || c == ',')
340                    .chars()
341                    .take(50)
342                    .collect();
343                if !title.is_empty() {
344                    return Some(title);
345                }
346            }
347            None
348        }))
349    }
350
351    // ===== Scrolling =====
352
353    pub fn scroll_up(&mut self, amount: u16) {
354        self.ui_state.chat_state.scroll_up(amount);
355    }
356
357    pub fn scroll_down(&mut self, amount: u16) {
358        self.ui_state.chat_state.scroll_down(amount);
359    }
360
361    // ===== Lifecycle =====
362
363    pub fn quit(&mut self) {
364        self.running = false;
365    }
366
367    // ===== Message History =====
368
369    /// Filter and prepare messages for model API calls.
370    /// Includes User, Assistant, and Tool messages for proper agent loop.
371    /// Injects timestamp context into User messages for the model's temporal awareness.
372    fn prepare_api_messages(&self) -> Vec<ChatMessage> {
373        self.session_state
374            .messages
375            .iter()
376            .filter(|msg| {
377                msg.role == MessageRole::User
378                    || msg.role == MessageRole::Assistant
379                    || msg.role == MessageRole::Tool
380            })
381            .map(|msg| {
382                if msg.role == MessageRole::User {
383                    let ts = msg.timestamp.format("%Y-%m-%d %H:%M:%S %Z").to_string();
384                    let mut m = msg.clone();
385                    m.content = format!("[Sent at: {}]\n{}", ts, m.content);
386                    m
387                } else {
388                    msg.clone()
389                }
390            })
391            .collect()
392    }
393
394    /// Build message history for model API calls (all messages, no truncation)
395    pub fn build_message_history(&self) -> Vec<ChatMessage> {
396        self.prepare_api_messages()
397    }
398
399    pub fn build_managed_message_history(
400        &self,
401        max_context_tokens: usize,
402        reserve_tokens: usize,
403    ) -> Vec<ChatMessage> {
404        use crate::utils::Tokenizer;
405
406        let tokenizer = Tokenizer::new(&self.model_state.model_name);
407        let available_tokens = max_context_tokens.saturating_sub(reserve_tokens);
408
409        let all_messages = self.prepare_api_messages();
410
411        if all_messages.is_empty() {
412            return Vec::new();
413        }
414
415        let messages_for_counting: Vec<(String, String)> = all_messages
416            .iter()
417            .map(|msg| {
418                let role = match msg.role {
419                    MessageRole::User => "user",
420                    MessageRole::Assistant => "assistant",
421                    MessageRole::System => "system",
422                    MessageRole::Tool => "tool",
423                };
424                (role.to_string(), msg.content.clone())
425            })
426            .collect();
427
428        let total_tokens = tokenizer
429            .count_chat_tokens(&messages_for_counting)
430            .unwrap_or_else(|_| all_messages.iter().map(|m| m.content.len() / 4).sum());
431
432        if total_tokens <= available_tokens {
433            return all_messages;
434        }
435
436        let mut kept_messages = Vec::new();
437        let mut current_tokens = 0;
438
439        for msg in all_messages.iter().rev() {
440            let msg_text = vec![(
441                match msg.role {
442                    MessageRole::User => "user",
443                    MessageRole::Assistant => "assistant",
444                    MessageRole::System => "system",
445                    MessageRole::Tool => "tool",
446                }
447                .to_string(),
448                msg.content.clone(),
449            )];
450
451            let msg_tokens = tokenizer
452                .count_chat_tokens(&msg_text)
453                .unwrap_or(msg.content.len() / 4);
454
455            if current_tokens + msg_tokens <= available_tokens {
456                kept_messages.push(msg.clone());
457                current_tokens += msg_tokens;
458            } else if kept_messages.len() < 2 {
459                kept_messages.push(msg.clone());
460                break;
461            } else {
462                break;
463            }
464        }
465
466        kept_messages.reverse();
467        kept_messages
468    }
469
470    // ===== Conversation Persistence =====
471
472    pub fn load_conversation(&mut self, conversation: ConversationHistory) {
473        self.session_state.cumulative_tokens = conversation.total_tokens.unwrap_or(0);
474        self.session_state.conversation_title = Some(conversation.title.clone());
475        self.session_state.messages = conversation.messages.clone();
476        self.session_state.current_conversation = Some(conversation);
477        self.set_status("Conversation loaded");
478    }
479
480    pub fn save_conversation(&mut self) -> anyhow::Result<()> {
481        if let Some(ref manager) = self.session_state.conversation_manager
482            && let Some(ref mut conv) = self.session_state.current_conversation
483        {
484            conv.messages = self.session_state.messages.clone();
485            conv.total_tokens = Some(self.session_state.cumulative_tokens);
486            manager.save_conversation(conv)?;
487            self.set_status("Conversation saved");
488        }
489        Ok(())
490    }
491
492    pub fn auto_save_conversation(&mut self) {
493        if self.session_state.messages.is_empty() {
494            return;
495        }
496        if let Some(ref manager) = self.session_state.conversation_manager
497            && let Some(ref mut conv) = self.session_state.current_conversation
498        {
499            conv.messages = self.session_state.messages.clone();
500            conv.total_tokens = Some(self.session_state.cumulative_tokens);
501            let conv_clone = conv.clone();
502            let manager_clone = manager.clone();
503            tokio::task::spawn_blocking(move || {
504                if let Err(e) = manager_clone.save_conversation(&conv_clone) {
505                    tracing::warn!("Failed to auto-save conversation: {}", e);
506                }
507            });
508        }
509    }
510
511    // ===== Generation State Transitions =====
512
513    pub fn start_generation(&mut self, abort_handle: tokio::task::AbortHandle) {
514        self.app_state = AppState::Generating {
515            status: GenerationStatus::Sending,
516            start_time: std::time::Instant::now(),
517            tokens_received: 0,
518            abort_handle: Some(abort_handle),
519            response_buffer: String::with_capacity(8192),
520        };
521    }
522
523    /// Update the abort handle for a new model call within the same turn.
524    /// Keeps the existing start_time and token count (cumulative for the turn).
525    pub fn update_abort_handle(&mut self, abort_handle: tokio::task::AbortHandle) {
526        if let AppState::Generating {
527            abort_handle: ref mut existing,
528            ..
529        } = self.app_state
530        {
531            *existing = Some(abort_handle);
532        }
533    }
534
535    /// Reset status to Sending for a new model call within the same turn.
536    pub fn transition_to_sending(&mut self) {
537        if let AppState::Generating { status, .. } = &mut self.app_state {
538            *status = GenerationStatus::Sending;
539        }
540    }
541
542    pub fn transition_to_thinking(&mut self) {
543        if let AppState::Generating { status, .. } = &mut self.app_state {
544            *status = GenerationStatus::Thinking;
545        }
546    }
547
548    pub fn transition_to_streaming(&mut self) {
549        if let AppState::Generating { status, .. } = &mut self.app_state {
550            *status = GenerationStatus::Streaming;
551        }
552    }
553
554    /// Add tokens from a completed model call (accumulates across the turn)
555    pub fn set_final_tokens(&mut self, count: usize) {
556        if let AppState::Generating {
557            tokens_received, ..
558        } = &mut self.app_state
559        {
560            *tokens_received += count;
561        }
562        self.session_state.add_tokens(count);
563    }
564
565    pub fn stop_generation(&mut self) {
566        self.app_state = AppState::Idle;
567    }
568
569    pub fn abort_generation(&mut self) -> (Option<tokio::task::AbortHandle>, String) {
570        if let AppState::Generating {
571            abort_handle,
572            response_buffer,
573            ..
574        } = &mut self.app_state
575        {
576            let handle = abort_handle.take();
577            let buffer = std::mem::take(response_buffer);
578            self.app_state = AppState::Idle;
579            (handle, buffer)
580        } else {
581            (None, String::new())
582        }
583    }
584
585    // ===== Response Buffer Accessors =====
586
587    /// Append text to the response buffer. No-op if not generating.
588    /// Enforces MAX_RESPONSE_CHARS size limit.
589    pub fn push_response(&mut self, text: &str) {
590        if let AppState::Generating {
591            response_buffer, ..
592        } = &mut self.app_state
593        {
594            response_buffer.push_str(text);
595            if response_buffer.len() > crate::constants::MAX_RESPONSE_CHARS {
596                let end = response_buffer.floor_char_boundary(crate::constants::MAX_RESPONSE_CHARS);
597                response_buffer.truncate(end);
598                response_buffer
599                    .push_str("\n\n[TRUNCATED: Response exceeded size limit]\n");
600                self.set_status("[WARNING] Response truncated (size limit reached)");
601            }
602        }
603    }
604
605    /// Get response buffer length (0 if not generating)
606    pub fn response_len(&self) -> usize {
607        if let AppState::Generating {
608            response_buffer, ..
609        } = &self.app_state
610        {
611            response_buffer.len()
612        } else {
613            0
614        }
615    }
616
617    /// Take the response buffer, leaving it empty. Returns empty string if not generating.
618    pub fn take_response(&mut self) -> String {
619        if let AppState::Generating {
620            response_buffer, ..
621        } = &mut self.app_state
622        {
623            std::mem::take(response_buffer)
624        } else {
625            String::new()
626        }
627    }
628
629    /// Clear the response buffer (for per-model-call reset within a turn)
630    pub fn clear_response(&mut self) {
631        if let AppState::Generating {
632            response_buffer, ..
633        } = &mut self.app_state
634        {
635            response_buffer.clear();
636        }
637    }
638}
mermaid_cli/tui/app.rs

mermaid_cli/tui/
app.rs
