elsa-core/test/integration/Elsa.AI.IntegrationTests/AIToolsEndpointTests.cs

using System.Reflection;
using System.Security.Claims;
using Elsa.AI.Abstractions.Contracts;
using Elsa.AI.Abstractions.Models;
using Elsa.AI.Host.Endpoints.AI.Tools;
using Elsa.AI.Host.Options;
using Elsa.Extensions;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.DependencyInjection;
using MicrosoftOptions = Microsoft.Extensions.Options.Options;
using Request = Elsa.AI.Host.Endpoints.AI.Tools.Request;
using ToolsEndpoint = Elsa.AI.Host.Endpoints.AI.Tools.Endpoint;

namespace Elsa.AI.IntegrationTests;

public class AIToolsEndpointTests
{
    [Fact(DisplayName = "Tools endpoint returns enabled registry results")]
    public async Task ToolsEndpointReturnsEnabledRegistryResults()
    {
        var services = new ServiceCollection();
        services.AddAIHostServices();
        using var provider = services.BuildServiceProvider();
        var endpoint = new ToolsEndpoint(provider.GetRequiredService<IAIToolRegistry>(), MicrosoftOptions.Create(new AIHostOptions()));

        var tools = await endpoint.ExecuteAsync(new Request(), CancellationToken.None);

        Assert.Contains(tools, x => x.Name == "activities.search");
        Assert.Contains(tools, x => x.Name == "workflows.search");
        Assert.Contains(tools, x => x.Name == "instances.search");
    }

    [Fact(DisplayName = "Tools endpoint forwards agent scope to registry")]
    public async Task ToolsEndpointForwardsAgentScopeToRegistry()
    {
        var services = new ServiceCollection();
        services.AddAIHostServices();
        services.AddSingleton<IAITool, WorkflowAuthorTool>();
        services.AddSingleton<IAITool, WorkflowEditorTool>();
        using var provider = services.BuildServiceProvider();
        var endpoint = new ToolsEndpoint(provider.GetRequiredService<IAIToolRegistry>(), MicrosoftOptions.Create(new AIHostOptions()));
        SetHttpContext(endpoint, "workflows:author");

        var tools = await endpoint.ExecuteAsync(new Request { Agent = "workflow-author" }, CancellationToken.None);

        Assert.Contains(tools, tool => tool.Name == "workflow.author");
        Assert.DoesNotContain(tools, tool => tool.Name == "workflow.editor");
    }

    [Fact(DisplayName = "Tool registry caches definitions across list calls")]
    public async Task ToolRegistryCachesDefinitionsAcrossListCalls()
    {
        CountingTool.Reset();
        var services = new ServiceCollection();
        services.AddAIHostServices();
        services.AddTransient<IAITool>(_ => CountingTool.Create());
        using var provider = services.BuildServiceProvider();
        var registry = provider.GetRequiredService<IAIToolRegistry>();

        await registry.ListAsync(new AIToolQuery(), CancellationToken.None);
        await registry.ListAsync(new AIToolQuery(), CancellationToken.None);

        Assert.Equal(1, CountingTool.ConstructorCount);
    }

    [Fact(DisplayName = "Tools endpoint lists built-in grounding tools")]
    public async Task ToolsEndpointListsBuiltInGroundingTools()
    {
        var services = new ServiceCollection();
        services.AddAIHostServices();
        using var provider = services.BuildServiceProvider();
        var endpoint = new ToolsEndpoint(provider.GetRequiredService<IAIToolRegistry>(), MicrosoftOptions.Create(new AIHostOptions()));

        var tools = await endpoint.ExecuteAsync(new Request(), CancellationToken.None);

        Assert.Contains(tools, tool => tool.Name == "activities.getDescriptor");
        Assert.Contains(tools, tool => tool.Name == "workflows.getDefinitionGraph");
        Assert.Contains(tools, tool => tool.Name == "workflows.validateDraft");
        Assert.Contains(tools, tool => tool.Name == "incidents.search");
        Assert.Contains(tools, tool => tool.Name == "workflows.proposeCreate" && !tool.IsEnabled);
    }

    private class WorkflowAuthorTool : IAITool
    {
        public AIToolDefinition Definition { get; } = new()
        {
            Name = "workflow.author",
            DisplayName = "Workflow author",
            AgentScopes = ["workflow-author"],
            Permissions = ["workflows:author"]
        };

        public ValueTask<AIToolResult> ExecuteAsync(AIToolExecutionContext context, CancellationToken cancellationToken = default) =>
            ValueTask.FromResult(new AIToolResult());

        public void Dispose()
        {
        }
    }

    private class WorkflowEditorTool : IAITool
    {
        public AIToolDefinition Definition { get; } = new()
        {
            Name = "workflow.editor",
            DisplayName = "Workflow editor",
            AgentScopes = ["workflow-editor"],
            Permissions = ["workflows:editor"]
        };

        public ValueTask<AIToolResult> ExecuteAsync(AIToolExecutionContext context, CancellationToken cancellationToken = default) =>
            ValueTask.FromResult(new AIToolResult());

        public void Dispose()
        {
        }
    }

    private class CountingTool : IAITool
    {
        private static int _constructorCount;

        public static int ConstructorCount => _constructorCount;

        private CountingTool()
        {
        }

        public static CountingTool Create()
        {
            Interlocked.Increment(ref _constructorCount);
            return new CountingTool();
        }

        public AIToolDefinition Definition { get; } = new()
        {
            Name = "counting.tool",
            DisplayName = "Counting tool"
        };

        public ValueTask<AIToolResult> ExecuteAsync(AIToolExecutionContext context, CancellationToken cancellationToken = default) =>
            ValueTask.FromResult(new AIToolResult());

        public void Dispose()
        {
        }

        public static void Reset()
        {
            Interlocked.Exchange(ref _constructorCount, 0);
        }
    }

    private static void SetHttpContext(ToolsEndpoint endpoint, params string[] permissions)
    {
        var identity = new ClaimsIdentity(permissions.Select(x => new Claim(PermissionNames.ClaimType, x)), "test");
        var property = typeof(ToolsEndpoint)
            .GetProperty(nameof(ToolsEndpoint.HttpContext), BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic)!;
        property.SetValue(endpoint, new DefaultHttpContext
        {
            User = new ClaimsPrincipal(identity)
        });
    }
}