Custom Tools

Create custom tools by implementing the Executable trait. Tools define their name, description, input schema, and execution logic.

The Executable Trait

The Executable trait is the foundation for all tools in agent-air. It defines the interface that tools must implement to be invoked by the LLM. The trait includes required methods for basic functionality and optional methods for advanced features like permissions and display customization.

pub trait Executable: Send + Sync {
    // Required methods
    fn name(&self) -> &str;
    fn description(&self) -> &str;
    fn input_schema(&self) -> &str;
    fn execute(
        &self,
        context: ToolContext,
        input: HashMap<String, serde_json::Value>,
    ) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>>;

    // Optional methods with defaults
    fn tool_type(&self) -> ToolType { ToolType::Custom }
    fn display_config(&self) -> DisplayConfig { DisplayConfig::default() }
    fn compact_summary(&self, input: &HashMap<String, Value>, result: &str) -> String;
    fn required_permissions(&self, input: &HashMap<String, Value>) -> Vec<PermissionRequest>;
    fn handles_own_permissions(&self) -> bool { false }
    fn cleanup_session(&self, session_id: i64) -> Pin<Box<dyn Future<Output = ()> + Send>>;
}

Required Methods

Every tool must implement these four methods. They define the tool’s identity, what it does, what inputs it accepts, and how it executes. The LLM uses the name, description, and input schema to decide when and how to call the tool.

name

Returns the tool’s unique identifier. The LLM uses this name to invoke the tool.

fn name(&self) -> &str {
    "calculate_sum"
}

Use lowercase with underscores. Keep names concise and descriptive.

description

Returns a description of what the tool does. The LLM uses this to decide when to use the tool.

fn description(&self) -> &str {
    "Calculates the sum of a list of numbers. Returns the total."
}

Be specific about inputs, outputs, and when the tool should be used.

input_schema

Returns a JSON Schema defining the tool’s input parameters.

fn input_schema(&self) -> &str {
    r#"{
        "type": "object",
        "properties": {
            "numbers": {
                "type": "array",
                "items": { "type": "number" },
                "description": "List of numbers to sum"
            }
        },
        "required": ["numbers"]
    }"#
}

See Input Schema Format for details.

execute

Performs the tool’s operation asynchronously.

fn execute(
    &self,
    context: ToolContext,
    input: HashMap<String, serde_json::Value>,
) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> {
    Box::pin(async move {
        let numbers = input.get("numbers")
            .and_then(|v| v.as_array())
            .ok_or("Missing numbers parameter")?;

        let sum: f64 = numbers.iter()
            .filter_map(|v| v.as_f64())
            .sum();

        Ok(format!("Sum: {}", sum))
    })
}

Optional Methods

These methods have default implementations but can be overridden to customize tool behavior. Use them to control how your tool appears in the UI, what permissions it requires, and how it handles session cleanup.

tool_type

Classifies the tool for permission and display purposes.

fn tool_type(&self) -> ToolType {
    ToolType::ApiCall
}

Available types:

display_config

Configures how the tool’s invocation appears in the UI.

fn display_config(&self) -> DisplayConfig {
    DisplayConfig {
        display_name: "Calculator".to_string(),
        display_title: Box::new(|input| {
            format!("Calculating sum of {} numbers",
                input.get("numbers")
                    .and_then(|v| v.as_array())
                    .map(|a| a.len())
                    .unwrap_or(0))
        }),
        display_content: Box::new(|_input, result| {
            DisplayResult {
                content: result.to_string(),
                content_type: ResultContentType::PlainText,
                is_truncated: false,
                full_length: result.len(),
            }
        }),
    }
}

compact_summary

Returns a short summary for context compaction.

fn compact_summary(
    &self,
    input: &HashMap<String, serde_json::Value>,
    result: &str,
) -> String {
    let count = input.get("numbers")
        .and_then(|v| v.as_array())
        .map(|a| a.len())
        .unwrap_or(0);
    format!("[Calculated sum of {} numbers: {}]", count, result)
}

required_permissions

Returns permissions needed before execution.

fn required_permissions(
    &self,
    input: &HashMap<String, serde_json::Value>,
) -> Vec<PermissionRequest> {
    let path = input.get("file_path")
        .and_then(|v| v.as_str())
        .unwrap_or("");

    vec![PermissionRequest {
        target_type: TargetType::Path,
        target: path.to_string(),
        level: PermissionLevel::Write,
        description: format!("Write to {}", path),
        reason: None,
        recursive: false,
    }]
}

handles_own_permissions

Return true if the tool manages its own permission flow.

fn handles_own_permissions(&self) -> bool {
    true // Tool will call PermissionRegistry directly
}

cleanup_session

Clean up session-specific state when a session ends.

fn cleanup_session(
    &self,
    session_id: i64,
) -> Pin<Box<dyn Future<Output = ()> + Send>> {
    let state = self.session_state.clone();
    Box::pin(async move {
        state.remove(&session_id);
    })
}

Input Schema Format

The input schema tells the LLM what parameters your tool accepts. Agent Air uses JSON Schema, a standard format that supports type validation, required fields, enums, and nested structures. A well-defined schema helps the LLM invoke your tool correctly.

Tools use JSON Schema to define their input parameters.

Basic Schema

{
    "type": "object",
    "properties": {
        "message": {
            "type": "string",
            "description": "The message to display"
        }
    },
    "required": ["message"]
}

Supported Types

JSON Schema supports several primitive and complex types. When extracting values in your execute method, use the corresponding serde_json methods to access the data.

Enums

Enums restrict a string parameter to a specific set of allowed values. The LLM will only provide values from this list.

{
    "type": "object",
    "properties": {
        "format": {
            "type": "string",
            "enum": ["json", "yaml", "toml"],
            "description": "Output format"
        }
    }
}

Nested Objects

Complex inputs can use nested objects to group related parameters logically.

{
    "type": "object",
    "properties": {
        "config": {
            "type": "object",
            "properties": {
                "host": { "type": "string" },
                "port": { "type": "integer" }
            }
        }
    }
}

Arrays with Typed Items

Arrays can specify the type of their elements using the items property.

{
    "type": "object",
    "properties": {
        "files": {
            "type": "array",
            "items": { "type": "string" },
            "description": "List of file paths"
        }
    }
}

Optional Parameters

Only include parameters in the required array if they are essential. Optional parameters can have default values that your tool applies when the LLM doesn’t provide them.

Parameters not in required are optional:

{
    "type": "object",
    "properties": {
        "query": { "type": "string" },
        "limit": { "type": "integer", "default": 10 }
    },
    "required": ["query"]
}

Content Types

Tool output can be rendered differently depending on its format. By specifying a content type, you help the UI choose appropriate syntax highlighting and formatting for your tool’s results.

The ResultContentType enum controls how tool output is rendered in the UI.

pub enum ResultContentType {
    Json,
    Markdown,
    Yaml,
    PlainText,
    Xml,
    Auto,
}

Using Content Types

Set the content type in your display configuration’s display_content function to ensure proper rendering.

fn display_config(&self) -> DisplayConfig {
    DisplayConfig {
        display_content: Box::new(|_input, result| {
            DisplayResult {
                content: result.to_string(),
                content_type: ResultContentType::Json,
                is_truncated: false,
                full_length: result.len(),
            }
        }),
        ..Default::default()
    }
}

Complete Example

This example demonstrates a complete tool implementation including all required methods and several optional ones. The weather tool makes an API call and returns formatted results with proper display configuration.

A tool that fetches weather data:

use agent_air::controller::{
    Executable, ToolContext, ToolType, DisplayConfig, DisplayResult,
    ResultContentType,
};
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;

pub struct WeatherTool {
    api_key: String,
}

impl WeatherTool {
    pub fn new(api_key: String) -> Self {
        Self { api_key }
    }
}

impl Executable for WeatherTool {
    fn name(&self) -> &str {
        "get_weather"
    }

    fn description(&self) -> &str {
        "Gets current weather for a city. Returns temperature, conditions, and humidity."
    }

    fn input_schema(&self) -> &str {
        r#"{
            "type": "object",
            "properties": {
                "city": {
                    "type": "string",
                    "description": "City name (e.g., 'San Francisco')"
                },
                "units": {
                    "type": "string",
                    "enum": ["celsius", "fahrenheit"],
                    "description": "Temperature units (default: celsius)"
                }
            },
            "required": ["city"]
        }"#
    }

    fn tool_type(&self) -> ToolType {
        ToolType::ApiCall
    }

    fn execute(
        &self,
        _context: 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 city = input.get("city")
                .and_then(|v| v.as_str())
                .ok_or("Missing city parameter")?;

            let units = input.get("units")
                .and_then(|v| v.as_str())
                .unwrap_or("celsius");

            // Make API call (simplified)
            let weather = fetch_weather(&api_key, city, units).await
                .map_err(|e| format!("API error: {}", e))?;

            Ok(serde_json::to_string_pretty(&weather).unwrap())
        })
    }

    fn display_config(&self) -> DisplayConfig {
        DisplayConfig {
            display_name: "Weather".to_string(),
            display_title: Box::new(|input| {
                input.get("city")
                    .and_then(|v| v.as_str())
                    .unwrap_or("Unknown")
                    .to_string()
            }),
            display_content: Box::new(|_input, result| {
                DisplayResult {
                    content: result.to_string(),
                    content_type: ResultContentType::Json,
                    is_truncated: false,
                    full_length: result.len(),
                }
            }),
        }
    }

    fn compact_summary(
        &self,
        input: &HashMap<String, serde_json::Value>,
        _result: &str,
    ) -> String {
        let city = input.get("city")
            .and_then(|v| v.as_str())
            .unwrap_or("unknown");
        format!("[Weather lookup: {}]", city)
    }
}

Registration

core.register_tools(|registry, _user_reg, _perm_reg| {
    let weather = Arc::new(WeatherTool::new(api_key));
    registry.register(weather.clone()).await?;
    Ok(vec![weather.to_llm_tool()])
})?;

Error Handling

Tools communicate failures by returning Err(String) from the execute method. The error message is passed back to the LLM, which can then inform the user or try an alternative approach. Keep error messages clear and actionable.

Return errors as Err(String) from execute():

fn execute(
    &self,
    _context: ToolContext,
    input: HashMap<String, serde_json::Value>,
) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> {
    Box::pin(async move {
        let path = input.get("path")
            .and_then(|v| v.as_str())
            .ok_or("Missing required parameter: path")?;

        if !path.starts_with('/') {
            return Err("Path must be absolute".to_string());
        }

        // Perform operation...
        Ok("Success".to_string())
    })
}

The error string is returned to the LLM, which can then decide how to handle the failure.