Agent Builder

The AgentAir builder provides methods for registering tools, widgets, and other components before running the agent. Registration happens during setup, before calling run(), and determines what capabilities are available during execution. Understanding these registration methods is key to extending agent functionality.

The builder pattern used by AgentAir allows fluent configuration through method chaining. Each registration method returns &mut Self, so you can configure multiple aspects in a single chain. Tools are registered with the LLM so it knows what operations are available, while widgets are registered with the TUI to extend the user interface.

Tool Registration

Tools are registered through the ToolRegistry and the register_tools() method on AgentAir. The registry provides thread-safe storage and lookup of tools, while the registration methods handle integration with LLM sessions.

ToolRegistry

The ToolRegistry is a thread-safe registry for managing available tools:

pub struct ToolRegistry {
    tools: RwLock<HashMap<String, Arc<dyn Executable>>>,
}

It uses a tokio RwLock to allow concurrent reads while ensuring exclusive writes.

Registry Methods

impl ToolRegistry {
    /// Create a new empty registry
    pub fn new() -> Self;

    /// Register a tool (returns error if name already exists)
    pub async fn register(&self, tool: Arc<dyn Executable>) -> Result<(), RegistryError>;

    /// Get a tool by name
    pub async fn get(&self, name: &str) -> Option<Arc<dyn Executable>>;

    /// Check if a tool exists
    pub async fn has(&self, name: &str) -> bool;

    /// List all tool names
    pub async fn list(&self) -> Vec<String>;

    /// Get all tools
    pub async fn get_all(&self) -> Vec<Arc<dyn Executable>>;

    /// Remove a tool by name
    pub async fn remove(&self, name: &str);

    /// Get the number of registered tools
    pub async fn len(&self) -> usize;

    /// Check if the registry is empty
    pub async fn is_empty(&self) -> bool;
}

Registering with AgentAir

Use register_tools() to register tools when building your agent:

impl AgentAir {
    pub fn register_tools<F>(&mut self, f: F) -> Result<(), AgentError>
    where
        F: FnOnce(
            &Arc<ToolRegistry>,
            &Arc<UserInteractionRegistry>,
            &Arc<PermissionRegistry>,
        ) -> Result<Vec<LLMTool>, String>;
}

The callback receives three registries:

• &Arc<ToolRegistry> - For registering Executable implementations
• &Arc<UserInteractionRegistry> - For tools that ask user questions
• &Arc<PermissionRegistry> - For tools that request permissions

The callback returns Vec<LLMTool> which are the tool definitions sent to the LLM.

Synchronous Registration

use agent_air::AgentAir;
use agent_air::controller::tools::{ToolRegistry, LLMTool};
use std::sync::Arc;

let mut agent = AgentAir::new(&MyConfig)?;

agent.register_tools(|registry, user_reg, perm_reg| {
    let my_tool = Arc::new(MyTool::new());

    // Register with the runtime
    tokio::runtime::Handle::current().block_on(async {
        registry.register(my_tool.clone()).await
    }).map_err(|e| e.to_string())?;

    // Return the LLM tool definition
    Ok(vec![my_tool.to_llm_tool()])
})?;

Async Registration

For tools that require async initialization, use register_tools_async():

impl AgentAir {
    pub fn register_tools_async<F, Fut>(&mut self, f: F) -> Result<(), AgentError>
    where
        F: FnOnce(
            Arc<ToolRegistry>,
            Arc<UserInteractionRegistry>,
            Arc<PermissionRegistry>
        ) -> Fut,
        Fut: Future<Output = Result<Vec<LLMTool>, String>>;
}

Example:

agent.register_tools_async(|registry, user_reg, perm_reg| async move {
    // Async initialization
    let client = HttpClient::connect("https://api.example.com").await?;

    let my_tool = Arc::new(ApiTool::new(client));
    registry.register(my_tool.clone()).await
        .map_err(|e| e.to_string())?;

    Ok(vec![my_tool.to_llm_tool()])
})?;

Tool Definition Flow

When tools are registered:

1. The callback registers Executable implementations with ToolRegistry
2. The callback returns Vec<LLMTool> definitions
3. AgentAir stores these definitions
4. When a session is created, tools are set on the session via session.set_tools()
5. The LLM receives tool definitions and can request their execution

Complete Tool Example

use agent_air::AgentAir;
use agent_air::controller::tools::{
    Executable, ToolContext, ToolType, LLMTool, ToolRegistry,
};
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

struct WeatherTool {
    api_key: String,
}

impl Executable for WeatherTool {
    fn name(&self) -> &str { "get_weather" }
    fn description(&self) -> &str { "Get weather for a location" }
    fn input_schema(&self) -> &str {
        r#"{"type":"object","properties":{"location":{"type":"string"}},"required":["location"]}"#
    }
    fn tool_type(&self) -> ToolType { ToolType::ApiCall }

    fn execute(
        &self,
        _ctx: ToolContext,
        input: HashMap<String, serde_json::Value>,
    ) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> {
        let api_key = self.api_key.clone();
        Box::pin(async move {
            let location = input.get("location")
                .and_then(|v| v.as_str())
                .ok_or("Missing location")?;
            Ok(format!("Weather in {}: Sunny", location))
        })
    }
}

fn main() -> std::io::Result<()> {
    let mut agent = AgentAir::new(&MyConfig)?;

    agent.register_tools_async(|registry, _user_reg, _perm_reg| async move {
        let weather = Arc::new(WeatherTool {
            api_key: std::env::var("WEATHER_API_KEY")
                .map_err(|_| "Missing WEATHER_API_KEY".to_string())?,
        });

        registry.register(weather.clone()).await
            .map_err(|e| e.to_string())?;

        Ok(vec![weather.to_llm_tool()])
    })?;

    agent.run()
}

RegistryError

Registration can fail with:

#[derive(Error, Debug)]
pub enum RegistryError {
    #[error("Tool with name {0:?} already exists")]
    DuplicateTool(String),
}

Accessing the Registry at Runtime

The registry is accessible through the controller:

let registry = agent.controller().tool_registry();

// Check if a tool exists
let has_weather = registry.has("get_weather").await;

// List all tools
let tool_names = registry.list().await;

// Remove a tool
registry.remove("deprecated_tool").await;

Widget Registration

Widgets are registered with AgentAir to add custom UI components to the TUI. Registered widgets participate in the rendering pipeline and can handle key events based on their priority. The widget system allows extending the interface with custom panels, overlays, and interactive elements.

The register_widget Method

Register widgets using register_widget() on AgentAir:

impl AgentAir {
    pub fn register_widget<W: Widget>(&mut self, widget: W) -> &mut Self;
}

This method returns &mut Self, allowing method chaining:

use agent_air::AgentAir;
use agent_air::tui::widgets::{PermissionPanel, QuestionPanel};

let mut agent = AgentAir::new(&MyConfig)?;

agent
    .register_widget(PermissionPanel::new())
    .register_widget(QuestionPanel::new());

agent.run()

Standard Widget IDs

The framework defines standard widget IDs:

pub mod widget_ids {
    // Core widgets (always present)
    pub const CHAT_VIEW: &str = "chat_view";
    pub const TEXT_INPUT: &str = "text_input";

    // Registerable widgets
    pub const PERMISSION_PANEL: &str = "permission_panel";
    pub const QUESTION_PANEL: &str = "question_panel";
    pub const SESSION_PICKER: &str = "session_picker";
    pub const SLASH_POPUP: &str = "slash_popup";
    pub const THEME_PICKER: &str = "theme_picker";
    pub const STATUS_BAR: &str = "status_bar";
}

Use these constants when implementing custom widgets or checking widget presence.

Widget Trait Requirements

Registered widgets must implement the Widget trait:

pub trait Widget: Send + 'static {
    fn id(&self) -> &'static str;
    fn priority(&self) -> u8 { 100 }
    fn is_active(&self) -> bool;
    fn handle_key(&mut self, key: KeyEvent, ctx: &WidgetKeyContext) -> WidgetKeyResult;
    fn render(&mut self, frame: &mut Frame, area: Rect, theme: &Theme);
    fn required_height(&self, available: u16) -> u16 { 0 }
    fn blocks_input(&self) -> bool { false }
    fn is_overlay(&self) -> bool { false }
    fn as_any(&self) -> &dyn Any;
    fn as_any_mut(&mut self) -> &mut dyn Any;
    fn into_any(self: Box<Self>) -> Box<dyn Any>;
}

Checking Widget Presence

The App provides methods to check if widgets are registered:

if app.has_widget(widget_ids::PERMISSION_PANEL) {
    // Permission panel is available
}

Accessing Registered Widgets

Access widgets by ID with type-safe downcasting:

use agent_air::tui::widgets::{widget_ids, PermissionPanel};

// Immutable access
if let Some(panel) = app.widget::<PermissionPanel>(widget_ids::PERMISSION_PANEL) {
    let is_active = panel.is_active();
}

// Mutable access
if let Some(panel) = app.widget_mut::<PermissionPanel>(widget_ids::PERMISSION_PANEL) {
    panel.set_options(options);
}

Priority-Based Key Handling

Widgets are checked for key handling in priority order (highest first). When a key event occurs:

1. App iterates widgets sorted by priority descending
2. Each active widget’s handle_key() is called
3. If a widget returns Handled or Action, processing stops
4. Otherwise, the next widget is checked

impl Widget for MyWidget {
    fn priority(&self) -> u8 {
        200  // Higher priority, checked before widgets with priority 100
    }

    fn handle_key(&mut self, key: KeyEvent, ctx: &WidgetKeyContext) -> WidgetKeyResult {
        if self.is_active() && key.code == KeyCode::Escape {
            self.close();
            return WidgetKeyResult::Handled;
        }
        WidgetKeyResult::NotHandled
    }
}

Overlay Widgets

Widgets that return true from is_overlay() are rendered on top of all other content:

impl Widget for FullScreenPicker {
    fn is_overlay(&self) -> bool {
        true  // Rendered last, on top of everything
    }
}

Blocking Input

Widgets that return true from blocks_input() prevent text input while active:

impl Widget for ModalDialog {
    fn blocks_input(&self) -> bool {
        self.is_active()  // Block text input when dialog is open
    }
}

Complete Widget Example

use agent_air::AgentAir;
use agent_air::tui::widgets::{
    Widget, WidgetKeyContext, WidgetKeyResult, widget_ids,
    PermissionPanel, QuestionPanel,
};
use crossterm::event::{KeyEvent, KeyCode};
use ratatui::{layout::Rect, Frame};
use std::any::Any;

struct NotificationWidget {
    message: Option<String>,
}

impl Widget for NotificationWidget {
    fn id(&self) -> &'static str { "notification" }
    fn priority(&self) -> u8 { 150 }
    fn is_active(&self) -> bool { self.message.is_some() }

    fn handle_key(&mut self, key: KeyEvent, _ctx: &WidgetKeyContext) -> WidgetKeyResult {
        if self.is_active() && matches!(key.code, KeyCode::Enter | KeyCode::Escape) {
            self.message = None;
            return WidgetKeyResult::Handled;
        }
        WidgetKeyResult::NotHandled
    }

    fn render(&mut self, frame: &mut Frame, area: Rect, theme: &Theme) {
        if let Some(msg) = &self.message {
            // Render notification...
        }
    }

    fn required_height(&self, _available: u16) -> u16 {
        if self.message.is_some() { 3 } else { 0 }
    }

    fn as_any(&self) -> &dyn Any { self }
    fn as_any_mut(&mut self) -> &mut dyn Any { self }
    fn into_any(self: Box<Self>) -> Box<dyn Any> { self }
}

fn main() -> std::io::Result<()> {
    let mut agent = AgentAir::new(&MyConfig)?;

    // Register built-in widgets
    agent
        .register_widget(PermissionPanel::new())
        .register_widget(QuestionPanel::new());

    // Register custom widget
    agent.register_widget(NotificationWidget { message: None });

    agent.run()
}

Widget Registration Order

Widgets are stored in registration order but processed by priority. Register widgets before calling run():

let mut agent = AgentAir::new(&config)?;

// All registrations must happen before run()
agent.register_widget(WidgetA::new());
agent.register_widget(WidgetB::new());
agent.register_widget(WidgetC::new());

// run() consumes the widget list
agent.run()

Builder Method Reference