1use 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
26pub type ToControllerTx = mpsc::Sender<ControllerInputPayload>;
28pub type ToControllerRx = mpsc::Receiver<ControllerInputPayload>;
30pub type FromControllerTx = mpsc::Sender<UiMessage>;
32pub type FromControllerRx = mpsc::Receiver<UiMessage>;
34
35pub struct AgentCore {
64 #[allow(dead_code)]
66 logger: Logger,
67
68 name: String,
70
71 version: String,
73
74 welcome_art: Vec<String>,
76
77 welcome_subtitle_indices: Vec<usize>,
79
80 runtime: Runtime,
82
83 controller: Arc<LLMController>,
85
86 llm_registry: Option<LLMRegistry>,
88
89 to_controller_tx: ToControllerTx,
91
92 to_controller_rx: Option<ToControllerRx>,
94
95 #[allow(dead_code)]
97 from_controller_tx: FromControllerTx,
98
99 from_controller_rx: Option<FromControllerRx>,
101
102 cancel_token: CancellationToken,
104
105 user_interaction_registry: Arc<UserInteractionRegistry>,
107
108 permission_registry: Arc<PermissionRegistry>,
110
111 tool_definitions: Vec<ToolDefinition>,
113
114 widgets_to_register: Vec<Box<dyn Widget>>,
116
117 layout_template: Option<LayoutTemplate>,
119
120 key_handler: Option<Box<dyn KeyHandler>>,
122
123 exit_handler: Option<Box<dyn ExitHandler>>,
125}
126
127impl AgentCore {
128 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 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 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 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 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 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 let (interaction_event_tx, mut interaction_event_rx) =
183 mpsc::channel::<ControllerEvent>(DEFAULT_CHANNEL_SIZE);
184
185 let user_interaction_registry =
187 Arc::new(UserInteractionRegistry::new(interaction_event_tx));
188
189 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 let (permission_event_tx, mut permission_event_rx) =
200 mpsc::channel::<ControllerEvent>(DEFAULT_CHANNEL_SIZE);
201
202 let permission_registry = Arc::new(PermissionRegistry::new(permission_event_tx));
204
205 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 pub fn set_version(&mut self, version: impl Into<String>) {
243 self.version = version.into();
244 }
245
246 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 pub fn set_layout(&mut self, template: LayoutTemplate) -> &mut Self {
270 self.layout_template = Some(template);
271 self
272 }
273
274 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 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 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 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 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 pub fn start_background_tasks(&mut self) {
396 tracing::info!("{} starting background tasks", self.name);
397
398 let controller = self.controller.clone();
400 self.runtime.spawn(async move {
401 controller.start().await;
402 });
403 tracing::info!("Controller started");
404
405 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 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 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 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 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 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 pub fn run(&mut self) -> io::Result<()> {
528 tracing::info!("{} starting", self.name);
529
530 self.start_background_tasks();
532
533 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 for widget in self.widgets_to_register.drain(..) {
545 let id = widget.id();
547 app.widgets.insert(id, widget);
548 }
549 app.rebuild_priority_order();
550
551 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 if let Some(layout) = self.layout_template.take() {
563 app.set_layout(layout);
564 }
565
566 if let Some(handler) = self.key_handler.take() {
568 app.set_key_handler_boxed(handler);
569 }
570
571 if let Some(handler) = self.exit_handler.take() {
573 app.set_exit_handler_boxed(handler);
574 }
575
576 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 if let Some(registry) = self.llm_registry.take() {
597 app.set_llm_registry(registry);
598 }
599
600 let result = app.run();
602
603 self.shutdown();
605
606 tracing::info!("{} stopped", self.name);
607 result
608 }
609
610 pub fn to_controller_tx(&self) -> ToControllerTx {
614 self.to_controller_tx.clone()
615 }
616
617 pub fn take_from_controller_rx(&mut self) -> Option<FromControllerRx> {
619 self.from_controller_rx.take()
620 }
621
622 pub fn controller(&self) -> &Arc<LLMController> {
624 &self.controller
625 }
626
627 pub fn runtime(&self) -> &Runtime {
629 &self.runtime
630 }
631
632 pub fn runtime_handle(&self) -> tokio::runtime::Handle {
634 self.runtime.handle().clone()
635 }
636
637 pub fn user_interaction_registry(&self) -> &Arc<UserInteractionRegistry> {
639 &self.user_interaction_registry
640 }
641
642 pub fn permission_registry(&self) -> &Arc<PermissionRegistry> {
644 &self.permission_registry
645 }
646
647 pub fn llm_registry(&self) -> Option<&LLMRegistry> {
649 self.llm_registry.as_ref()
650 }
651
652 pub fn take_llm_registry(&mut self) -> Option<LLMRegistry> {
654 self.llm_registry.take()
655 }
656
657 pub fn cancel_token(&self) -> CancellationToken {
659 self.cancel_token.clone()
660 }
661
662 pub fn name(&self) -> &str {
664 &self.name
665 }
666}
667
668pub fn convert_controller_event_to_ui_message(event: ControllerEvent) -> UiMessage {
673 match event {
674 ControllerEvent::StreamStart { session_id, .. } => {
675 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}