gitdataai/libs/agent/tool/registry.rs

//! Request-scoped tool registry.
//!
//! Tools are registered per-request (not globally) to keep the system testable
//! and allow different request contexts to have different tool sets.

use std::collections::HashMap;

use futures::FutureExt;

use super::call::ToolError;
use super::context::ToolContext;
use super::definition::ToolDefinition;

/// Inner function pointer type for tool handlers.
type InnerHandlerFn = dyn Fn(
        ToolContext,
        serde_json::Value,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = Result<serde_json::Value, ToolError>> + Send>,
    > + Send
    + Sync;

/// Wrapper around `Arc<dyn Fn(...)>` for `Clone` implementability.
#[derive(Clone)]
pub struct ToolHandler(std::sync::Arc<InnerHandlerFn>);

impl ToolHandler {
    /// Creates a new handler from an async closure.
    /// The closure should return `Result<serde_json::Value, String>` (as used by git_tools),
    /// which is converted to `Result<serde_json::Value, ToolError>`.
    pub fn new<F>(f: F) -> Self
    where
        F: Fn(ToolContext, serde_json::Value) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<serde_json::Value, ToolError>> + Send>>
            + Send
            + Sync
            + 'static,
    {
        Self(std::sync::Arc::new(f))
    }

    pub async fn execute(
        &self,
        ctx: ToolContext,
        args: serde_json::Value,
    ) -> Result<serde_json::Value, ToolError> {
        (self.0)(ctx, args).await
    }
}

/// A request-scoped registry mapping tool names to their handlers.
#[derive(Clone, Default)]
pub struct ToolRegistry {
    handlers: HashMap<String, ToolHandler>,
    definitions: HashMap<String, ToolDefinition>,
}

impl ToolRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn register_fn<F, Fut>(&mut self, name: impl Into<String>, handler: F) -> &mut Self
    where
        F: Fn(ToolContext, serde_json::Value) -> Fut + Send + Sync + 'static,
        Fut: std::future::Future<Output = Result<serde_json::Value, ToolError>> + Send + 'static,
    {
        let name_str = name.into();
        let def = ToolDefinition::new(&name_str);
        let handler_fn: std::sync::Arc<InnerHandlerFn> =
            std::sync::Arc::new(move |ctx, args| handler(ctx, args).boxed());
        self.register(def, ToolHandler(handler_fn));
        self
    }

    pub fn register(&mut self, def: ToolDefinition, handler: ToolHandler) -> &mut Self {
        let name = def.name.clone();
        if self.handlers.contains_key(&name) {
            panic!("tool already registered: {}", name);
        }
        self.handlers.insert(name.clone(), handler);
        self.definitions.insert(name, def);
        self
    }

    /// Looks up a handler by tool name.
    pub fn get(&self, name: &str) -> Option<&ToolHandler> {
        self.handlers.get(name)
    }

    pub fn definitions(&self) -> std::collections::hash_map::Values<'_, String, ToolDefinition> {
        self.definitions.values()
    }

    pub fn to_openai_tools(&self) -> Vec<async_openai::types::chat::ChatCompletionTool> {
        self.definitions
            .values()
            .map(|d| d.to_openai_tool())
            .collect()
    }

    pub fn len(&self) -> usize {
        self.handlers.len()
    }

    pub fn is_empty(&self) -> bool {
        self.handlers.is_empty()
    }

    /// Merges another registry's tools into this one.
    /// Panics if a tool with the same name already exists.
    pub fn merge(&mut self, other: ToolRegistry) {
        for (name, handler) in other.handlers {
            if self.handlers.contains_key(&name) {
                panic!("tool already registered: {}", name);
            }
            self.handlers.insert(name, handler);
        }
        for (name, def) in other.definitions {
            self.definitions.insert(name, def);
        }
    }
}