agent_core/agent/
core.rs

1// AgentCore - Complete working agent out of the box
2//
3// Users call `AgentCore::new(config)` then `core.run()` and get a working chat agent.
4
5use std::io;
6use std::sync::Arc;
7
8use tokio::runtime::Runtime;
9use tokio::sync::mpsc;
10use tokio_util::sync::CancellationToken;
11
12use crate::controller::{
13    ControllerEvent, ControllerInputPayload, LLMController, LLMSessionConfig, LLMTool,
14    PermissionRegistry, ToolDefinition, ToolRegistry, UserInteractionRegistry,
15};
16
17use super::config::{load_config, AgentConfig, LLMRegistry};
18use super::logger::Logger;
19use super::messages::channels::DEFAULT_CHANNEL_SIZE;
20use super::messages::UiMessage;
21use super::router::InputRouter;
22
23use crate::tui::{App, AppConfig, DefaultKeyHandler, ExitHandler, KeyBindings, KeyHandler, LayoutTemplate, SessionInfo};
24use crate::tui::widgets::Widget;
25
26/// Sender for messages from TUI to controller
27pub type ToControllerTx = mpsc::Sender<ControllerInputPayload>;
28/// Receiver for messages from TUI to controller
29pub type ToControllerRx = mpsc::Receiver<ControllerInputPayload>;
30/// Sender for messages from controller to TUI
31pub type FromControllerTx = mpsc::Sender<UiMessage>;
32/// Receiver for messages from controller to TUI
33pub type FromControllerRx = mpsc::Receiver<UiMessage>;
34
35/// AgentCore - A complete, working agent infrastructure.
36///
37/// AgentCore provides all the infrastructure needed for an LLM-powered agent:
38/// - Logging with tracing
39/// - LLM configuration loading
40/// - Tokio async runtime
41/// - LLMController for session management
42/// - Communication channels
43/// - User interaction and permission registries
44///
45/// # Basic Usage
46///
47/// ```ignore
48/// struct MyConfig;
49/// impl AgentConfig for MyConfig {
50///     fn config_path(&self) -> &str { ".myagent/config.yaml" }
51///     fn default_system_prompt(&self) -> &str { "You are helpful." }
52///     fn log_prefix(&self) -> &str { "myagent" }
53///     fn name(&self) -> &str { "MyAgent" }
54/// }
55///
56/// fn main() -> io::Result<()> {
57///     let mut core = AgentCore::new(&MyConfig)?;
58///     // Access channels and controller to wire up your TUI
59///     // then run your TUI loop
60///     Ok(())
61/// }
62/// ```
63pub struct AgentCore {
64    /// Logger instance (must be kept alive)
65    #[allow(dead_code)]
66    logger: Logger,
67
68    /// Agent name for display
69    name: String,
70
71    /// Agent version for display
72    version: String,
73
74    /// Welcome ASCII art lines
75    welcome_art: Vec<String>,
76
77    /// Subtitle line indices in welcome art
78    welcome_subtitle_indices: Vec<usize>,
79
80    /// Tokio runtime for async operations
81    runtime: Runtime,
82
83    /// The LLM controller
84    controller: Arc<LLMController>,
85
86    /// LLM provider registry (loaded from config)
87    llm_registry: Option<LLMRegistry>,
88
89    /// Sender for messages from TUI to controller
90    to_controller_tx: ToControllerTx,
91
92    /// Receiver for messages from TUI to controller (consumed by InputRouter)
93    to_controller_rx: Option<ToControllerRx>,
94
95    /// Sender for messages from controller to TUI (held by event handler)
96    #[allow(dead_code)]
97    from_controller_tx: FromControllerTx,
98
99    /// Receiver for messages from controller to TUI
100    from_controller_rx: Option<FromControllerRx>,
101
102    /// Cancellation token for graceful shutdown
103    cancel_token: CancellationToken,
104
105    /// User interaction registry for AskUserQuestions tool
106    user_interaction_registry: Arc<UserInteractionRegistry>,
107
108    /// Permission registry for AskForPermissions tool
109    permission_registry: Arc<PermissionRegistry>,
110
111    /// Tool definitions to register on sessions
112    tool_definitions: Vec<ToolDefinition>,
113
114    /// Widgets to register with the App
115    widgets_to_register: Vec<Box<dyn Widget>>,
116
117    /// Layout template for the TUI
118    layout_template: Option<LayoutTemplate>,
119
120    /// Key handler for customizable key bindings
121    key_handler: Option<Box<dyn KeyHandler>>,
122
123    /// Exit handler for cleanup before quitting
124    exit_handler: Option<Box<dyn ExitHandler>>,
125}
126
127impl AgentCore {
128    /// Create a new AgentCore with the given configuration.
129    ///
130    /// This initializes:
131    /// - Logging infrastructure
132    /// - LLM configuration from config file or environment
133    /// - Tokio runtime
134    /// - Communication channels
135    /// - LLMController
136    /// - User interaction and permission registries
137    pub fn new<C: AgentConfig>(config: &C) -> io::Result<Self> {
138        let logger = Logger::new(config.log_prefix())?;
139        tracing::info!("{} agent initialized", config.name());
140
141        // Load LLM configuration
142        let llm_registry = load_config(config);
143        if llm_registry.is_empty() {
144            tracing::warn!(
145                "No LLM providers configured. Set ANTHROPIC_API_KEY or create ~/{}",
146                config.config_path()
147            );
148        } else {
149            tracing::info!(
150                "Loaded {} LLM provider(s): {:?}",
151                llm_registry.providers().len(),
152                llm_registry.providers()
153            );
154        }
155
156        // Create tokio runtime for async operations
157        let runtime = Runtime::new().map_err(|e| {
158            io::Error::new(
159                io::ErrorKind::Other,
160                format!("Failed to create runtime: {}", e),
161            )
162        })?;
163
164        // Create communication channels
165        let (to_controller_tx, to_controller_rx) =
166            mpsc::channel::<ControllerInputPayload>(DEFAULT_CHANNEL_SIZE);
167        let (from_controller_tx, from_controller_rx) =
168            mpsc::channel::<UiMessage>(DEFAULT_CHANNEL_SIZE);
169
170        // Create the controller with an event handler that forwards to the UI channel
171        let ui_tx = from_controller_tx.clone();
172        let event_handler = Box::new(move |event: ControllerEvent| {
173            let msg = convert_controller_event_to_ui_message(event);
174            // Try to send, but don't block if channel is full
175            let _ = ui_tx.try_send(msg);
176        });
177
178        let controller = Arc::new(LLMController::new(Some(event_handler)));
179        let cancel_token = CancellationToken::new();
180
181        // Create channel for user interaction events
182        let (interaction_event_tx, mut interaction_event_rx) =
183            mpsc::channel::<ControllerEvent>(DEFAULT_CHANNEL_SIZE);
184
185        // Create the user interaction registry
186        let user_interaction_registry =
187            Arc::new(UserInteractionRegistry::new(interaction_event_tx));
188
189        // Spawn a task to forward user interaction events to the UI channel
190        let ui_tx_for_interactions = from_controller_tx.clone();
191        runtime.spawn(async move {
192            while let Some(event) = interaction_event_rx.recv().await {
193                let msg = convert_controller_event_to_ui_message(event);
194                let _ = ui_tx_for_interactions.try_send(msg);
195            }
196        });
197
198        // Create channel for permission events
199        let (permission_event_tx, mut permission_event_rx) =
200            mpsc::channel::<ControllerEvent>(DEFAULT_CHANNEL_SIZE);
201
202        // Create the permission registry
203        let permission_registry = Arc::new(PermissionRegistry::new(permission_event_tx));
204
205        // Spawn a task to forward permission events to the UI channel
206        let ui_tx_for_permissions = from_controller_tx.clone();
207        runtime.spawn(async move {
208            while let Some(event) = permission_event_rx.recv().await {
209                let msg = convert_controller_event_to_ui_message(event);
210                let _ = ui_tx_for_permissions.try_send(msg);
211            }
212        });
213
214        Ok(Self {
215            logger,
216            name: config.name().to_string(),
217            version: "0.1.0".to_string(),
218            welcome_art: vec![
219                String::new(),
220                "    Type a message to start chatting...".to_string(),
221            ],
222            welcome_subtitle_indices: vec![1],
223            runtime,
224            controller,
225            llm_registry: Some(llm_registry),
226            to_controller_tx,
227            to_controller_rx: Some(to_controller_rx),
228            from_controller_tx,
229            from_controller_rx: Some(from_controller_rx),
230            cancel_token,
231            user_interaction_registry,
232            permission_registry,
233            tool_definitions: Vec::new(),
234            widgets_to_register: Vec::new(),
235            layout_template: None,
236            key_handler: None,
237            exit_handler: None,
238        })
239    }
240
241    /// Set the agent version for display.
242    pub fn set_version(&mut self, version: impl Into<String>) {
243        self.version = version.into();
244    }
245
246    /// Set the welcome ASCII art displayed when chat is empty.
247    pub fn set_welcome_art(&mut self, art: Vec<String>, subtitle_indices: Vec<usize>) {
248        self.welcome_art = art;
249        self.welcome_subtitle_indices = subtitle_indices;
250    }
251
252    /// Set the layout template for the TUI.
253    ///
254    /// This allows customizing how widgets are arranged in the terminal.
255    /// If not set, the default Standard layout with panels is used.
256    ///
257    /// # Example
258    ///
259    /// ```ignore
260    /// // Use standard layout (default)
261    /// agent.set_layout(LayoutTemplate::standard());
262    ///
263    /// // Add a sidebar
264    /// agent.set_layout(LayoutTemplate::with_sidebar("file_browser", 30));
265    ///
266    /// // Minimal layout (no status bar)
267    /// agent.set_layout(LayoutTemplate::minimal());
268    /// ```
269    pub fn set_layout(&mut self, template: LayoutTemplate) -> &mut Self {
270        self.layout_template = Some(template);
271        self
272    }
273
274    /// Set a custom key handler for the TUI.
275    ///
276    /// This allows full control over key handling behavior. For simpler
277    /// customization where you just want to change which keys trigger
278    /// which actions, use [`set_key_bindings`] instead.
279    ///
280    /// # Example
281    ///
282    /// ```ignore
283    /// struct VimKeyHandler { mode: VimMode }
284    /// impl KeyHandler for VimKeyHandler {
285    ///     fn handle_key(&mut self, key: KeyEvent, ctx: &KeyContext) -> AppKeyResult {
286    ///         // Implement vim-style modal editing
287    ///     }
288    /// }
289    ///
290    /// let mut agent = AgentCore::new(&config)?;
291    /// agent.set_key_handler(VimKeyHandler { mode: VimMode::Normal });
292    /// ```
293    pub fn set_key_handler<H: KeyHandler>(&mut self, handler: H) -> &mut Self {
294        self.key_handler = Some(Box::new(handler));
295        self
296    }
297
298    /// Set custom key bindings using the default handler.
299    ///
300    /// This is a simpler alternative to [`set_key_handler`] when you
301    /// only need to change which keys trigger which actions.
302    ///
303    /// # Example
304    ///
305    /// ```ignore
306    /// // Use minimal bindings (Esc to quit, arrow keys only)
307    /// let mut agent = AgentCore::new(&config)?;
308    /// agent.set_key_bindings(KeyBindings::minimal());
309    ///
310    /// // Or customize specific bindings
311    /// let mut bindings = KeyBindings::emacs();
312    /// bindings.quit = vec![KeyCombo::key(KeyCode::Esc)];
313    /// agent.set_key_bindings(bindings);
314    /// ```
315    pub fn set_key_bindings(&mut self, bindings: KeyBindings) -> &mut Self {
316        self.key_handler = Some(Box::new(DefaultKeyHandler::new(bindings)));
317        self
318    }
319
320    /// Set an exit handler for cleanup before quitting.
321    ///
322    /// The exit handler's `on_exit()` method is called when the user
323    /// confirms exit. If it returns `false`, the exit is cancelled.
324    ///
325    /// # Example
326    ///
327    /// ```ignore
328    /// struct SaveOnExitHandler { session_file: PathBuf }
329    /// impl ExitHandler for SaveOnExitHandler {
330    ///     fn on_exit(&mut self) -> bool {
331    ///         self.save_session();
332    ///         true // proceed with exit
333    ///     }
334    /// }
335    ///
336    /// let mut agent = AgentCore::new(&config)?;
337    /// agent.set_exit_handler(SaveOnExitHandler { session_file: path });
338    /// ```
339    pub fn set_exit_handler<H: ExitHandler>(&mut self, handler: H) -> &mut Self {
340        self.exit_handler = Some(Box::new(handler));
341        self
342    }
343
344    /// Register tools with the agent.
345    ///
346    /// The callback receives references to the tool registry and interaction registries,
347    /// and should return the tool definitions to register.
348    ///
349    /// # Example
350    ///
351    /// ```ignore
352    /// core.register_tools(|registry, user_reg, perm_reg| {
353    ///     tools::register_all_tools(registry, user_reg, perm_reg)
354    /// })?;
355    /// ```
356    pub fn register_tools<F>(&mut self, f: F) -> Result<(), String>
357    where
358        F: FnOnce(
359            &Arc<ToolRegistry>,
360            &Arc<UserInteractionRegistry>,
361            &Arc<PermissionRegistry>,
362        ) -> Result<Vec<ToolDefinition>, String>,
363    {
364        let tool_defs = f(
365            self.controller.tool_registry(),
366            &self.user_interaction_registry,
367            &self.permission_registry,
368        )?;
369        self.tool_definitions = tool_defs;
370        Ok(())
371    }
372
373    /// Register a widget with the agent.
374    ///
375    /// Widgets are registered before calling `run()` and will be available
376    /// in the TUI application.
377    ///
378    /// # Example
379    ///
380    /// ```ignore
381    /// let mut agent = AgentCore::new(&MyConfig)?;
382    /// agent.register_widget(PermissionPanel::new());
383    /// agent.register_widget(QuestionPanel::new());
384    /// agent.run()
385    /// ```
386    pub fn register_widget<W: Widget>(&mut self, widget: W) -> &mut Self {
387        self.widgets_to_register.push(Box::new(widget));
388        self
389    }
390
391    /// Start the controller and input router as background tasks.
392    ///
393    /// This must be called before sending messages or creating sessions.
394    /// After calling this, the controller is running and ready to accept input.
395    pub fn start_background_tasks(&mut self) {
396        tracing::info!("{} starting background tasks", self.name);
397
398        // Start the controller event loop in a background task
399        let controller = self.controller.clone();
400        self.runtime.spawn(async move {
401            controller.start().await;
402        });
403        tracing::info!("Controller started");
404
405        // Start the input router in a background task
406        if let Some(to_controller_rx) = self.to_controller_rx.take() {
407            let router = InputRouter::new(
408                self.controller.clone(),
409                to_controller_rx,
410                self.cancel_token.clone(),
411            );
412            self.runtime.spawn(async move {
413                router.run().await;
414            });
415            tracing::info!("InputRouter started");
416        }
417    }
418
419    /// Create an initial session using the default LLM provider.
420    ///
421    /// Returns the session ID and model name, or an error message.
422    pub fn create_initial_session(&mut self) -> Result<(i64, String, i32), String> {
423        let registry = self.llm_registry.as_ref().ok_or_else(|| {
424            "No LLM registry available. Configuration may have failed to load.".to_string()
425        })?;
426
427        let config = registry
428            .get_default()
429            .ok_or_else(|| "No default LLM provider configured.".to_string())?;
430
431        let model = config.model.clone();
432        let context_limit = config.context_limit;
433        let session_config = config.clone();
434
435        let controller = self.controller.clone();
436        let tool_definitions = self.tool_definitions.clone();
437
438        let session_id = self
439            .runtime
440            .block_on(async {
441                let id = controller.create_session(session_config).await?;
442
443                // Set tools on the session after creation
444                if !tool_definitions.is_empty() {
445                    let tools: Vec<LLMTool> = tool_definitions
446                        .iter()
447                        .map(|def| LLMTool::new(&def.name, &def.description, &def.input_schema))
448                        .collect();
449
450                    if let Some(session) = controller.get_session(id).await {
451                        session.set_tools(tools).await;
452                    }
453                }
454
455                Ok::<i64, crate::client::error::LlmError>(id)
456            })
457            .map_err(|e| format!("Failed to create session: {}", e))?;
458
459        tracing::info!(
460            session_id = session_id,
461            model = %model,
462            "Created initial session"
463        );
464
465        Ok((session_id, model, context_limit))
466    }
467
468    /// Create a session with the given configuration.
469    ///
470    /// Returns the session ID or an error.
471    pub fn create_session(&self, config: LLMSessionConfig) -> Result<i64, String> {
472        let controller = self.controller.clone();
473        let tool_definitions = self.tool_definitions.clone();
474
475        self.runtime
476            .block_on(async {
477                let id = controller.create_session(config).await?;
478
479                // Set tools on the session after creation
480                if !tool_definitions.is_empty() {
481                    let tools: Vec<LLMTool> = tool_definitions
482                        .iter()
483                        .map(|def| LLMTool::new(&def.name, &def.description, &def.input_schema))
484                        .collect();
485
486                    if let Some(session) = controller.get_session(id).await {
487                        session.set_tools(tools).await;
488                    }
489                }
490
491                Ok::<i64, crate::client::error::LlmError>(id)
492            })
493            .map_err(|e| format!("Failed to create session: {}", e))
494    }
495
496    /// Signal shutdown to all background tasks and the controller.
497    pub fn shutdown(&self) {
498        tracing::info!("{} shutting down", self.name);
499        self.cancel_token.cancel();
500
501        let controller = self.controller.clone();
502        self.runtime.block_on(async move {
503            controller.shutdown().await;
504        });
505
506        tracing::info!("{} shutdown complete", self.name);
507    }
508
509    /// Run the agent with the default TUI.
510    ///
511    /// This is the main entry point for running an agent. It:
512    /// 1. Starts background tasks (controller, input router)
513    /// 2. Creates an App with the configured settings
514    /// 3. Wires up all channels and registries
515    /// 4. Creates an initial session if LLM providers are configured
516    /// 5. Runs the TUI event loop
517    /// 6. Shuts down cleanly when the user quits
518    ///
519    /// # Example
520    ///
521    /// ```ignore
522    /// fn main() -> io::Result<()> {
523    ///     let mut agent = AgentCore::new(&MyConfig)?;
524    ///     agent.run()
525    /// }
526    /// ```
527    pub fn run(&mut self) -> io::Result<()> {
528        tracing::info!("{} starting", self.name);
529
530        // Start background tasks (controller, input router)
531        self.start_background_tasks();
532
533        // Create App with our configuration
534        let app_config = AppConfig {
535            agent_name: self.name.clone(),
536            version: self.version.clone(),
537            welcome_art: self.welcome_art.clone(),
538            welcome_subtitle_indices: self.welcome_subtitle_indices.clone(),
539            custom_commands: Vec::new(),
540        };
541        let mut app = App::with_config(app_config);
542
543        // Register widgets with the App
544        for widget in self.widgets_to_register.drain(..) {
545            // We need to re-box as the App's register_widget expects impl Widget
546            let id = widget.id();
547            app.widgets.insert(id, widget);
548        }
549        app.rebuild_priority_order();
550
551        // Wire up channels, controller, and registries to the App
552        app.set_to_controller(self.to_controller_tx.clone());
553        if let Some(rx) = self.from_controller_rx.take() {
554            app.set_from_controller(rx);
555        }
556        app.set_controller(self.controller.clone());
557        app.set_runtime_handle(self.runtime.handle().clone());
558        app.set_user_interaction_registry(self.user_interaction_registry.clone());
559        app.set_permission_registry(self.permission_registry.clone());
560
561        // Set layout template if specified
562        if let Some(layout) = self.layout_template.take() {
563            app.set_layout(layout);
564        }
565
566        // Set key handler if specified
567        if let Some(handler) = self.key_handler.take() {
568            app.set_key_handler_boxed(handler);
569        }
570
571        // Set exit handler if specified
572        if let Some(handler) = self.exit_handler.take() {
573            app.set_exit_handler_boxed(handler);
574        }
575
576        // Auto-create session if we have a configured LLM provider
577        match self.create_initial_session() {
578            Ok((session_id, model, context_limit)) => {
579                let session_info = SessionInfo::new(session_id, model.clone(), context_limit);
580                app.add_session(session_info);
581                app.set_session_id(session_id);
582                app.set_model_name(&model);
583                app.set_context_limit(context_limit);
584                tracing::info!(
585                    session_id = session_id,
586                    model = %model,
587                    "Auto-created session on startup"
588                );
589            }
590            Err(e) => {
591                tracing::warn!(error = %e, "No initial session created");
592            }
593        }
594
595        // Pass LLM registry to app for creating new sessions
596        if let Some(registry) = self.llm_registry.take() {
597            app.set_llm_registry(registry);
598        }
599
600        // Run the TUI (blocking)
601        let result = app.run();
602
603        // Shutdown the agent
604        self.shutdown();
605
606        tracing::info!("{} stopped", self.name);
607        result
608    }
609
610    // ---- Accessors ----
611
612    /// Returns a sender for sending messages to the controller.
613    pub fn to_controller_tx(&self) -> ToControllerTx {
614        self.to_controller_tx.clone()
615    }
616
617    /// Takes the receiver for messages from the controller (can only be called once).
618    pub fn take_from_controller_rx(&mut self) -> Option<FromControllerRx> {
619        self.from_controller_rx.take()
620    }
621
622    /// Returns a reference to the controller.
623    pub fn controller(&self) -> &Arc<LLMController> {
624        &self.controller
625    }
626
627    /// Returns a reference to the runtime.
628    pub fn runtime(&self) -> &Runtime {
629        &self.runtime
630    }
631
632    /// Returns a handle to the runtime.
633    pub fn runtime_handle(&self) -> tokio::runtime::Handle {
634        self.runtime.handle().clone()
635    }
636
637    /// Returns a reference to the user interaction registry.
638    pub fn user_interaction_registry(&self) -> &Arc<UserInteractionRegistry> {
639        &self.user_interaction_registry
640    }
641
642    /// Returns a reference to the permission registry.
643    pub fn permission_registry(&self) -> &Arc<PermissionRegistry> {
644        &self.permission_registry
645    }
646
647    /// Returns a reference to the LLM registry.
648    pub fn llm_registry(&self) -> Option<&LLMRegistry> {
649        self.llm_registry.as_ref()
650    }
651
652    /// Takes the LLM registry (can only be called once).
653    pub fn take_llm_registry(&mut self) -> Option<LLMRegistry> {
654        self.llm_registry.take()
655    }
656
657    /// Returns the cancellation token.
658    pub fn cancel_token(&self) -> CancellationToken {
659        self.cancel_token.clone()
660    }
661
662    /// Returns the agent name.
663    pub fn name(&self) -> &str {
664        &self.name
665    }
666}
667
668/// Converts a ControllerEvent to a UiMessage for the TUI.
669///
670/// This function maps the internal controller events to UI-friendly messages
671/// that can be displayed in a terminal interface.
672pub fn convert_controller_event_to_ui_message(event: ControllerEvent) -> UiMessage {
673    match event {
674        ControllerEvent::StreamStart { session_id, .. } => {
675            // Silent - don't display stream start messages
676            UiMessage::System {
677                session_id,
678                message: String::new(),
679            }
680        }
681        ControllerEvent::TextChunk {
682            session_id,
683            text,
684            turn_id,
685        } => UiMessage::TextChunk {
686            session_id,
687            turn_id,
688            text,
689            input_tokens: 0,
690            output_tokens: 0,
691        },
692        ControllerEvent::ToolUseStart {
693            session_id,
694            tool_name,
695            turn_id,
696            ..
697        } => UiMessage::Display {
698            session_id,
699            turn_id,
700            message: format!("Executing tool: {}", tool_name),
701        },
702        ControllerEvent::ToolUse {
703            session_id,
704            tool,
705            display_name,
706            display_title,
707            turn_id,
708        } => UiMessage::ToolExecuting {
709            session_id,
710            turn_id,
711            tool_use_id: tool.id.clone(),
712            display_name: display_name.unwrap_or_else(|| tool.name.clone()),
713            display_title: display_title.unwrap_or_default(),
714        },
715        ControllerEvent::Complete {
716            session_id,
717            turn_id,
718            stop_reason,
719        } => UiMessage::Complete {
720            session_id,
721            turn_id,
722            input_tokens: 0,
723            output_tokens: 0,
724            stop_reason,
725        },
726        ControllerEvent::Error {
727            session_id,
728            error,
729            turn_id,
730        } => UiMessage::Error {
731            session_id,
732            turn_id,
733            error,
734        },
735        ControllerEvent::TokenUpdate {
736            session_id,
737            input_tokens,
738            output_tokens,
739            context_limit,
740        } => UiMessage::TokenUpdate {
741            session_id,
742            turn_id: None,
743            input_tokens,
744            output_tokens,
745            context_limit,
746        },
747        ControllerEvent::ToolResult {
748            session_id,
749            tool_use_id,
750            status,
751            error,
752            turn_id,
753            ..
754        } => UiMessage::ToolCompleted {
755            session_id,
756            turn_id,
757            tool_use_id,
758            status,
759            error,
760        },
761        ControllerEvent::CommandComplete {
762            session_id,
763            command,
764            success,
765            message,
766        } => UiMessage::CommandComplete {
767            session_id,
768            command,
769            success,
770            message,
771        },
772        ControllerEvent::UserInteractionRequired {
773            session_id,
774            tool_use_id,
775            request,
776            turn_id,
777        } => UiMessage::UserInteractionRequired {
778            session_id,
779            tool_use_id,
780            request,
781            turn_id,
782        },
783        ControllerEvent::PermissionRequired {
784            session_id,
785            tool_use_id,
786            request,
787            turn_id,
788        } => UiMessage::PermissionRequired {
789            session_id,
790            tool_use_id,
791            request,
792            turn_id,
793        },
794    }
795}
796
797#[cfg(test)]
798mod tests {
799    use super::*;
800    use crate::controller::TurnId;
801
802    struct TestConfig;
803
804    impl AgentConfig for TestConfig {
805        fn config_path(&self) -> &str {
806            ".test_agent/config.yaml"
807        }
808
809        fn default_system_prompt(&self) -> &str {
810            "You are a test agent."
811        }
812
813        fn log_prefix(&self) -> &str {
814            "test_agent"
815        }
816
817        fn name(&self) -> &str {
818            "TestAgent"
819        }
820    }
821
822    #[test]
823    fn test_convert_text_chunk_event() {
824        let event = ControllerEvent::TextChunk {
825            session_id: 1,
826            text: "Hello".to_string(),
827            turn_id: Some(TurnId::new_user_turn(1)),
828        };
829
830        let msg = convert_controller_event_to_ui_message(event);
831
832        match msg {
833            UiMessage::TextChunk {
834                session_id, text, ..
835            } => {
836                assert_eq!(session_id, 1);
837                assert_eq!(text, "Hello");
838            }
839            _ => panic!("Expected TextChunk message"),
840        }
841    }
842
843    #[test]
844    fn test_convert_error_event() {
845        let event = ControllerEvent::Error {
846            session_id: 1,
847            error: "Test error".to_string(),
848            turn_id: None,
849        };
850
851        let msg = convert_controller_event_to_ui_message(event);
852
853        match msg {
854            UiMessage::Error {
855                session_id, error, ..
856            } => {
857                assert_eq!(session_id, 1);
858                assert_eq!(error, "Test error");
859            }
860            _ => panic!("Expected Error message"),
861        }
862    }
863}