Skip to main content

agent_air_tui/
runner.rs

1//! TUI Runner - Provides the run() method for AgentAir
2//!
3//! This module extends AgentAir with TUI-specific functionality,
4//! including widget registration, layout configuration, and the main
5//! TUI event loop.
6
7use std::io;
8
9use agent_air_runtime::agent::AgentAir;
10
11use crate::app::{App, AppConfig};
12use crate::commands::SlashCommand;
13use crate::keys::{DefaultKeyHandler, ExitHandler, KeyBindings, KeyHandler};
14use crate::layout::LayoutTemplate;
15use crate::widgets::{ConversationView, ConversationViewFactory, SessionInfo, Widget, widget_ids};
16
17/// TUI configuration that extends AgentAir with TUI-specific settings.
18///
19/// This struct holds all TUI-related configuration that would otherwise
20/// need to be stored on AgentAir. Use `TuiRunner::new()` to create one
21/// from an AgentAir, configure it, then call `run()`.
22pub struct TuiRunner {
23    /// The underlying agent air
24    agent: AgentAir,
25
26    /// Factory for creating conversation views
27    conversation_factory: Option<ConversationViewFactory>,
28
29    /// Widgets to register with the App
30    widgets_to_register: Vec<Box<dyn Widget>>,
31
32    /// Layout template for the TUI
33    layout_template: Option<LayoutTemplate>,
34
35    /// Key handler for customizable key bindings
36    key_handler: Option<Box<dyn KeyHandler>>,
37
38    /// Exit handler for cleanup before quitting
39    exit_handler: Option<Box<dyn ExitHandler>>,
40
41    /// Slash commands (None means use defaults)
42    commands: Option<Vec<Box<dyn SlashCommand>>>,
43
44    /// Extension data available to commands
45    command_extension: Option<Box<dyn std::any::Any + Send>>,
46
47    /// Custom status bar widget (replaces default if provided)
48    custom_status_bar: Option<Box<dyn Widget>>,
49
50    /// Whether to hide the default status bar
51    hide_status_bar: bool,
52}
53
54impl TuiRunner {
55    /// Create a new TuiRunner from an AgentAir.
56    pub fn new(agent: AgentAir) -> Self {
57        Self {
58            agent,
59            conversation_factory: None,
60            widgets_to_register: Vec::new(),
61            layout_template: None,
62            key_handler: None,
63            exit_handler: None,
64            commands: None,
65            command_extension: None,
66            custom_status_bar: None,
67            hide_status_bar: false,
68        }
69    }
70
71    /// Get a mutable reference to the underlying AgentAir.
72    pub fn agent_mut(&mut self) -> &mut AgentAir {
73        &mut self.agent
74    }
75
76    /// Get a reference to the underlying AgentAir.
77    pub fn agent(&self) -> &AgentAir {
78        &self.agent
79    }
80
81    /// Set the conversation view factory.
82    ///
83    /// The factory is called to create conversation views when sessions
84    /// are created or cleared. This allows customizing the chat view
85    /// with custom welcome screens, title renderers, etc.
86    pub fn set_conversation_factory<F>(&mut self, factory: F) -> &mut Self
87    where
88        F: Fn() -> Box<dyn ConversationView> + Send + Sync + 'static,
89    {
90        self.conversation_factory = Some(Box::new(factory));
91        self
92    }
93
94    /// Set the layout template for the TUI.
95    ///
96    /// This allows customizing how widgets are arranged in the terminal.
97    /// If not set, the default Standard layout with panels is used.
98    pub fn set_layout(&mut self, template: LayoutTemplate) -> &mut Self {
99        self.layout_template = Some(template);
100        self
101    }
102
103    /// Set a custom key handler for the TUI.
104    ///
105    /// This allows full control over key handling behavior.
106    pub fn set_key_handler<H: KeyHandler>(&mut self, handler: H) -> &mut Self {
107        self.key_handler = Some(Box::new(handler));
108        self
109    }
110
111    /// Set custom key bindings using the default handler.
112    ///
113    /// This is a simpler alternative to `set_key_handler` when you
114    /// only need to change which keys trigger which actions.
115    pub fn set_key_bindings(&mut self, bindings: KeyBindings) -> &mut Self {
116        self.key_handler = Some(Box::new(DefaultKeyHandler::new(bindings)));
117        self
118    }
119
120    /// Set an exit handler for cleanup before quitting.
121    ///
122    /// The exit handler's `on_exit()` method is called when the user
123    /// confirms exit. If it returns `false`, the exit is cancelled.
124    pub fn set_exit_handler<H: ExitHandler>(&mut self, handler: H) -> &mut Self {
125        self.exit_handler = Some(Box::new(handler));
126        self
127    }
128
129    /// Set the slash commands for this agent.
130    ///
131    /// If not called, uses the default command set.
132    pub fn set_commands(&mut self, commands: Vec<Box<dyn SlashCommand>>) -> &mut Self {
133        self.commands = Some(commands);
134        self
135    }
136
137    /// Set extension data available to custom commands.
138    ///
139    /// Commands can access this via `ctx.extension::<T>()`.
140    pub fn set_command_extension<T: std::any::Any + Send + 'static>(
141        &mut self,
142        ext: T,
143    ) -> &mut Self {
144        self.command_extension = Some(Box::new(ext));
145        self
146    }
147
148    /// Register a widget with the agent.
149    ///
150    /// Widgets are registered before calling `run()` and will be available
151    /// in the TUI application.
152    pub fn register_widget<W: Widget>(&mut self, widget: W) -> &mut Self {
153        self.widgets_to_register.push(Box::new(widget));
154        self
155    }
156
157    /// Set a custom status bar widget to replace the default.
158    pub fn set_status_bar<W: Widget>(&mut self, status_bar: W) -> &mut Self {
159        self.custom_status_bar = Some(Box::new(status_bar));
160        self
161    }
162
163    /// Hide the default status bar.
164    pub fn hide_status_bar(&mut self) -> &mut Self {
165        self.hide_status_bar = true;
166        self
167    }
168
169    /// Run the agent with the TUI.
170    ///
171    /// This is the main entry point for running an agent with TUI. It:
172    /// 1. Starts background tasks (controller, input router)
173    /// 2. Creates an App with the configured settings
174    /// 3. Wires up all channels and registries
175    /// 4. Creates an initial session if LLM providers are configured
176    /// 5. Runs the TUI event loop
177    /// 6. Shuts down cleanly when the user quits
178    pub fn run(mut self) -> io::Result<()> {
179        let name = self.agent.name().to_string();
180        tracing::info!("{} starting", name);
181
182        // Start background tasks (controller, input router)
183        self.agent.start_background_tasks();
184
185        // Create App with our configuration
186        let app_config = AppConfig {
187            agent_name: name.clone(),
188            version: self.agent.version().to_string(),
189            commands: self.commands.take(),
190            command_extension: self.command_extension.take(),
191            error_no_session: self.agent.error_no_session().map(|s| s.to_string()),
192            ..Default::default()
193        };
194        let mut app = App::with_config(app_config);
195
196        // Set conversation factory if provided
197        if let Some(factory) = self.conversation_factory.take() {
198            app.set_conversation_factory(move || factory());
199        }
200
201        // Handle status bar customization
202        if self.hide_status_bar {
203            // Remove the default status bar
204            app.widgets.remove(widget_ids::STATUS_BAR);
205        } else if let Some(custom_status_bar) = self.custom_status_bar.take() {
206            // Replace default status bar with custom one
207            app.widgets
208                .insert(widget_ids::STATUS_BAR, custom_status_bar);
209        }
210
211        // Register widgets with the App
212        for widget in self.widgets_to_register.drain(..) {
213            let id = widget.id();
214            app.widgets.insert(id, widget);
215        }
216        app.rebuild_priority_order();
217
218        // Wire up channels, controller, and registries to the App
219        app.set_to_controller(self.agent.to_controller_tx());
220        if let Some(rx) = self.agent.take_from_controller_rx() {
221            app.set_from_controller(rx);
222        }
223        app.set_controller(self.agent.controller().clone());
224        app.set_runtime_handle(self.agent.runtime_handle());
225        app.set_user_interaction_registry(self.agent.user_interaction_registry().clone());
226        app.set_permission_registry(self.agent.permission_registry().clone());
227
228        // Set layout template if specified
229        if let Some(layout) = self.layout_template.take() {
230            app.set_layout(layout);
231        }
232
233        // Set key handler if specified
234        if let Some(handler) = self.key_handler.take() {
235            app.set_key_handler_boxed(handler);
236        }
237
238        // Set exit handler if specified
239        if let Some(handler) = self.exit_handler.take() {
240            app.set_exit_handler_boxed(handler);
241        }
242
243        // Auto-create session if we have a configured LLM provider
244        match self.agent.create_initial_session() {
245            Ok((session_id, model, context_limit)) => {
246                let session_info = SessionInfo::new(session_id, model.clone(), context_limit);
247                app.add_session(session_info);
248                app.set_session_id(session_id);
249                app.set_model_name(&model);
250                app.set_context_limit(context_limit);
251                tracing::info!(
252                    session_id = session_id,
253                    model = %model,
254                    "Auto-created session on startup"
255                );
256            }
257            Err(e) => {
258                tracing::warn!(error = %e, "No initial session created");
259            }
260        }
261
262        // Pass LLM registry to app for creating new sessions
263        if let Some(registry) = self.agent.take_llm_registry() {
264            app.set_llm_registry(registry);
265        }
266
267        // Run the TUI (blocking)
268        let result = app.run();
269
270        // Shutdown the agent
271        self.agent.shutdown();
272
273        tracing::info!("{} stopped", name);
274        result
275    }
276}
277
278/// Extension trait to add TUI functionality to AgentAir.
279///
280/// This trait is implemented for AgentAir when the `tui` feature is enabled,
281/// providing a convenient `into_tui()` method.
282pub trait AgentAirExt {
283    /// Convert this AgentAir into a TuiRunner for TUI operation.
284    fn into_tui(self) -> TuiRunner;
285}
286
287impl AgentAirExt for AgentAir {
288    fn into_tui(self) -> TuiRunner {
289        TuiRunner::new(self)
290    }
291}