golem-rust 3.0.0-rc2

// Copyright 2024-2026 Golem Cloud
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::agentic::InputStream;
use crate::agentic::extended_tool_type::{ExtendedToolType, PreparedToolDescriptor};
use crate::golem_agentic::exports::golem::tool::guest::{
    InvocationResult, Tool, ToolError, TypedSchemaValue,
};
use crate::golem_agentic::golem::agent::common::Principal;
use crate::golem_agentic::golem::tool::streams::ToolOutputWriter;
use std::cell::RefCell;
use std::collections::BTreeMap;
use std::future::Future;
use std::pin::Pin;

pub type ToolInvokeFuture =
    Pin<Box<dyn Future<Output = Result<InvocationResult, ToolError>> + 'static>>;

pub type ToolInvokeFutureFor<'a> =
    Pin<Box<dyn Future<Output = Result<InvocationResult, ToolError>> + 'a>>;

pub type ToolInvoker = fn(
    Vec<String>,
    TypedSchemaValue,
    Option<InputStream>,
    Option<ToolOutputWriter>,
    Option<ToolOutputWriter>,
    Principal,
) -> ToolInvokeFuture;

#[derive(Default)]
pub struct State {
    pub tools: RefCell<Tools>,
}
#[derive(Default)]
pub struct Tools {
    pub tools: BTreeMap<String, Box<dyn Fn() -> Tool>>,
    pub invokers: BTreeMap<String, ToolInvoker>,
}
static mut STATE: Option<State> = None;

#[allow(static_mut_refs)]
pub fn get_state() -> &'static State {
    unsafe {
        if STATE.is_none() {
            STATE = Some(State::default());
        }
        STATE.as_ref().unwrap()
    }
}

pub fn register_tool(tool: ExtendedToolType) {
    register_tool_inner(tool.prepare().expect("tool descriptor build failed"), None)
}

pub fn register_tool_invoker(tool: ExtendedToolType, invoker: ToolInvoker) {
    register_prepared_tool_invoker(
        tool.prepare().expect("tool descriptor build failed"),
        invoker,
    )
}

#[doc(hidden)]
pub fn register_prepared_tool_invoker(tool: PreparedToolDescriptor, invoker: ToolInvoker) {
    register_tool_inner(tool, Some(invoker))
}

fn register_tool_inner(tool: PreparedToolDescriptor, invoker: Option<ToolInvoker>) {
    let name = tool.extended().tool_name().to_string();
    register_factory(name, Box::new(move || tool.wire()), invoker);
}

#[doc(hidden)]
pub fn register_wire_tool_invoker(name: &str, descriptor: fn() -> Tool, invoker: ToolInvoker) {
    register_factory(name.to_string(), Box::new(descriptor), Some(invoker));
}

fn register_factory(name: String, descriptor: Box<dyn Fn() -> Tool>, invoker: Option<ToolInvoker>) {
    let state = get_state();
    let mut tools = state.tools.borrow_mut();
    if tools.tools.contains_key(&name) {
        panic!("duplicate tool registration for tool name: {name}");
    }
    if let Some(invoker) = invoker {
        tools.invokers.insert(name.clone(), invoker);
    }
    tools.tools.insert(name, descriptor);
}

pub fn get_all_tools() -> Vec<Tool> {
    get_state()
        .tools
        .borrow()
        .tools
        .values()
        .map(|tool| tool())
        .collect()
}
pub fn get_tool_by_name(name: &str) -> Option<Tool> {
    get_state()
        .tools
        .borrow()
        .tools
        .get(name)
        .map(|tool| tool())
}
pub fn get_extended_tool_by_name(name: &str) -> Option<ExtendedToolType> {
    get_tool_by_name(name).map(|tool| {
        crate::schema::tool::wit::decode_tool(tool)
            .expect("registered descriptor must be valid")
            .into()
    })
}

pub fn get_tool_invoker_by_name(name: &str) -> Option<ToolInvoker> {
    get_state().tools.borrow().invokers.get(name).copied()
}

#[cfg(test)]
pub(crate) fn clear_tools_for_tests() {
    let mut tools = get_state().tools.borrow_mut();
    tools.tools.clear();
    tools.invokers.clear();
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agentic::{Doc, ExtendedCommandNode, ExtendedGlobals, ExtendedToolType};
    use test_r::test;

    fn tool(name: &str) -> ExtendedToolType {
        ExtendedToolType {
            version: "0.1.0".into(),
            requires_filesystem: false,
            commands: vec![ExtendedCommandNode {
                name: name.into(),
                aliases: vec![],
                doc: Doc {
                    summary: String::new(),
                    description: String::new(),
                    examples: vec![],
                },
                globals: ExtendedGlobals::default(),
                subcommands: vec![],
                body: None,
            }],
        }
    }

    #[test]
    #[should_panic]
    fn duplicate_registration_panics() {
        clear_tools_for_tests();
        register_tool(tool("dupe"));
        register_tool(tool("dupe"));
    }
}