1use 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
17pub struct TuiRunner {
23 agent: AgentAir,
25
26 conversation_factory: Option<ConversationViewFactory>,
28
29 widgets_to_register: Vec<Box<dyn Widget>>,
31
32 layout_template: Option<LayoutTemplate>,
34
35 key_handler: Option<Box<dyn KeyHandler>>,
37
38 exit_handler: Option<Box<dyn ExitHandler>>,
40
41 commands: Option<Vec<Box<dyn SlashCommand>>>,
43
44 command_extension: Option<Box<dyn std::any::Any + Send>>,
46
47 custom_status_bar: Option<Box<dyn Widget>>,
49
50 hide_status_bar: bool,
52}
53
54impl TuiRunner {
55 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 pub fn agent_mut(&mut self) -> &mut AgentAir {
73 &mut self.agent
74 }
75
76 pub fn agent(&self) -> &AgentAir {
78 &self.agent
79 }
80
81 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 pub fn set_layout(&mut self, template: LayoutTemplate) -> &mut Self {
99 self.layout_template = Some(template);
100 self
101 }
102
103 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 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 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 pub fn set_commands(&mut self, commands: Vec<Box<dyn SlashCommand>>) -> &mut Self {
133 self.commands = Some(commands);
134 self
135 }
136
137 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 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 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 pub fn hide_status_bar(&mut self) -> &mut Self {
165 self.hide_status_bar = true;
166 self
167 }
168
169 pub fn run(mut self) -> io::Result<()> {
179 let name = self.agent.name().to_string();
180 tracing::info!("{} starting", name);
181
182 self.agent.start_background_tasks();
184
185 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 if let Some(factory) = self.conversation_factory.take() {
198 app.set_conversation_factory(move || factory());
199 }
200
201 if self.hide_status_bar {
203 app.widgets.remove(widget_ids::STATUS_BAR);
205 } else if let Some(custom_status_bar) = self.custom_status_bar.take() {
206 app.widgets
208 .insert(widget_ids::STATUS_BAR, custom_status_bar);
209 }
210
211 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 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 if let Some(layout) = self.layout_template.take() {
230 app.set_layout(layout);
231 }
232
233 if let Some(handler) = self.key_handler.take() {
235 app.set_key_handler_boxed(handler);
236 }
237
238 if let Some(handler) = self.exit_handler.take() {
240 app.set_exit_handler_boxed(handler);
241 }
242
243 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 if let Some(registry) = self.agent.take_llm_registry() {
264 app.set_llm_registry(registry);
265 }
266
267 let result = app.run();
269
270 self.agent.shutdown();
272
273 tracing::info!("{} stopped", name);
274 result
275 }
276}
277
278pub trait AgentAirExt {
283 fn into_tui(self) -> TuiRunner;
285}
286
287impl AgentAirExt for AgentAir {
288 fn into_tui(self) -> TuiRunner {
289 TuiRunner::new(self)
290 }
291}