Tool Registration & Context

Tools must be registered with the agent before they can be used. During execution, tools receive context about the current session and request.

Tool Context

Every tool invocation receives contextual information about the current session and request. This context enables tools to maintain session-specific state, track invocations for logging, and understand their position in the conversation.

The ToolContext struct provides execution context to tools:

pub struct ToolContext {
    pub session_id: i64,
    pub tool_use_id: String,
    pub turn_id: i64,
}

Using Context

The context is passed automatically by the framework when invoking a tool. Access its fields to implement session-aware behavior or correlate tool calls with their results.

Tools receive context as the first parameter to 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 {
        // Access context fields
        let session = context.session_id;
        let invocation = context.tool_use_id;

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

Context Use Cases

Different context fields serve different purposes depending on your tool’s requirements. Here are common patterns for using each field.

Session-scoped state: Use session_id to maintain per-session state like working directories or cached permissions.

Request tracking: Use tool_use_id to correlate tool invocations with their results, useful for logging and debugging.

Conversation position: Use turn_id to understand where in the conversation this tool is being called.

Tool Registry

The tool registry is a centralized store for all tools available to the agent. It handles registration, lookup, and lifecycle management of tools. Tools must be registered before the LLM can invoke them.

The ToolRegistry manages registered tools and provides lookup functionality.

Registering Tools

Tool registration happens during agent initialization through a closure that receives the necessary registries. This pattern ensures all dependencies are available when tools are created.

Use the register_tools method on AgentAir:

core.register_tools(|registry, user_reg, perm_reg| {
    // Register a tool
    let my_tool = Arc::new(MyTool::new());
    registry.register(my_tool.clone()).await?;

    // Return the list of tools to expose to the LLM
    Ok(vec![my_tool.to_llm_tool()])
})?;

The closure receives three registries:

Registry Methods

The registry provides methods for adding, retrieving, and querying tools. These methods are async and thread-safe, allowing concurrent access from multiple tasks.

register: Add a tool to the registry.

registry.register(Arc::new(MyTool::new())).await?;

get: Retrieve a tool by name.

if let Some(tool) = registry.get("my_tool").await {
    // Use tool
}

list: Get all registered tool names.

let tool_names = registry.list().await;

contains: Check if a tool is registered.

if registry.contains("my_tool").await {
    // Tool exists
}

User Interaction Registry

Tools that need to ask users questions use the user interaction registry to coordinate the request-response flow. The registry handles the async communication between tools and the UI, blocking tool execution until the user responds.

The UserInteractionRegistry manages pending user questions and their responses.

Purpose

When a tool needs to ask the user a question, it:

1. Registers the question with UserInteractionRegistry
2. Receives a channel to wait on
3. Blocks until the user responds
4. Receives the response through the channel

Usage

pub struct MyInteractiveTool {
    registry: Arc<UserInteractionRegistry>,
}

impl MyInteractiveTool {
    pub fn new(registry: Arc<UserInteractionRegistry>) -> Self {
        Self { registry }
    }
}

In the execute method:

fn execute(
    &self,
    context: ToolContext,
    input: HashMap<String, serde_json::Value>,
) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> {
    let registry = self.registry.clone();

    Box::pin(async move {
        // Create the question request
        let request = UserQuestionRequest { /* ... */ };

        // Register and wait for response
        let rx = registry.register(
            context.tool_use_id,
            context.session_id,
            request,
            context.turn_id,
        ).await?;

        // Block until user responds
        let response = rx.await.map_err(|e| e.to_string())?;

        Ok(serde_json::to_string(&response).unwrap())
    })
}

Permission Registry

Tools that perform sensitive operations use the permission registry to request user approval. The registry handles permission prompts and caches session-level grants so users don’t need to approve the same permission repeatedly.

The PermissionRegistry manages permission requests and caches session-level grants.

Purpose

When a tool needs permission for a sensitive operation, it:

1. Checks if permission is already granted for this session
2. If not, registers a permission request
3. Waits for user response
4. Caches session-level grants for future requests

Usage

pub struct MyPermissionTool {
    registry: Arc<PermissionRegistry>,
}

impl MyPermissionTool {
    pub fn new(registry: Arc<PermissionRegistry>) -> Self {
        Self { registry }
    }
}

In the execute method:

fn execute(
    &self,
    context: ToolContext,
    input: HashMap<String, serde_json::Value>,
) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> {
    let registry = self.registry.clone();

    Box::pin(async move {
        let request = PermissionRequest { /* ... */ };

        // Check if already granted for this session
        if registry.is_granted(context.session_id, &request).await {
            // Permission already granted, proceed
            return Ok("Operation completed".to_string());
        }

        // Register permission request and wait
        let rx = registry.register(
            context.tool_use_id,
            context.session_id,
            request,
            context.turn_id,
        ).await?;

        let response = rx.await.map_err(|e| e.to_string())?;

        if response.granted {
            // User granted permission, proceed
            Ok("Operation completed".to_string())
        } else {
            Err("Permission denied".to_string())
        }
    })
}

Session Grants

Session-scoped permissions reduce user friction by remembering approvals for similar operations. Once granted, permissions of the same type are automatically approved without prompting.

When a user grants permission with “session” scope, subsequent requests of the same type are automatically approved:

// First request prompts user
registry.register(tool_use_id, session_id, request, turn_id).await?;

// Later, same type of request
if registry.is_granted(session_id, &similar_request).await {
    // Automatically approved, no prompt needed
}

Registration Patterns

Different tools have different dependency requirements. These patterns show how to register tools based on whether they need user interaction, permissions, or neither.

Basic Registration

Register a simple tool without dependencies:

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

With User Interaction

Register a tool that asks questions:

core.register_tools(|registry, user_reg, _perm_reg| {
    let tool = Arc::new(InteractiveTool::new(user_reg.clone()));
    registry.register(tool.clone()).await?;
    Ok(vec![tool.to_llm_tool()])
})?;

With Permissions

Register a tool that requires permissions:

core.register_tools(|registry, _user_reg, perm_reg| {
    let tool = Arc::new(SensitiveTool::new(perm_reg.clone()));
    registry.register(tool.clone()).await?;
    Ok(vec![tool.to_llm_tool()])
})?;

Multiple Tools

Register multiple tools at once:

core.register_tools(|registry, user_reg, perm_reg| {
    let read_tool = Arc::new(ReadFileTool::new());
    let write_tool = Arc::new(WriteFileTool::new(perm_reg.clone()));
    let ask_tool = Arc::new(AskUserQuestionsTool::new(user_reg.clone()));

    registry.register(read_tool.clone()).await?;
    registry.register(write_tool.clone()).await?;
    registry.register(ask_tool.clone()).await?;

    Ok(vec![
        read_tool.to_llm_tool(),
        write_tool.to_llm_tool(),
        ask_tool.to_llm_tool(),
    ])
})?;

LLM Tool Conversion

After registering a tool, you must also return it in the format the LLM expects. This conversion extracts the tool’s name, description, and input schema into a structure the LLM provider can understand.

The to_llm_tool() method converts an Executable into the format expected by the LLM:

pub fn to_llm_tool(&self) -> LLMTool {
    LLMTool {
        name: self.name().to_string(),
        description: self.description().to_string(),
        input_schema: serde_json::from_str(self.input_schema()).unwrap(),
    }
}

The returned LLMTool is sent to the LLM provider so the model knows what tools are available.

Session Cleanup

Tools that maintain session-specific state should clean up when sessions end. The framework automatically calls the cleanup method on all registered tools, ensuring resources are released properly.

Tools can implement cleanup logic that runs when a session ends:

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

The framework calls cleanup_session on all registered tools when a session terminates, allowing tools to release resources or clear cached state.