using Elsa.AI.Abstractions.Contracts; using Elsa.AI.Abstractions.Models; using Elsa.AI.Host.Services; using Microsoft.Extensions.DependencyInjection; namespace Elsa.AI.Host.UnitTests; public class AIToolRegistryTests { [Fact(DisplayName = "Tool registry excludes host and cross-tenant denied tools from tenant queries")] public async Task ToolRegistryExcludesHostAndCrossTenantDeniedToolsFromTenantQueries() { var registry = CreateRegistry( [ new TestTool(new AIToolDefinition { Name = "tenant", DisplayName = "Tenant", TenantBehavior = AITenantBehavior.TenantScoped }), new TestTool(new AIToolDefinition { Name = "host", DisplayName = "Host", TenantBehavior = AITenantBehavior.HostScoped }), new TestTool(new AIToolDefinition { Name = "cross", DisplayName = "Cross", TenantBehavior = AITenantBehavior.CrossTenantDenied }) ]); var tools = await registry.ListAsync(new AIToolQuery { TenantId = "tenant-1", ActorId = "user-1" }); var tool = Assert.Single(tools); Assert.Equal("tenant", tool.Name); } [Fact(DisplayName = "Tool registry treats empty tenant ID as tenant context")] public async Task ToolRegistryTreatsEmptyTenantIdAsTenantContext() { var registry = CreateRegistry( [ new TestTool(new AIToolDefinition { Name = "tenant", DisplayName = "Tenant", TenantBehavior = AITenantBehavior.TenantScoped }), new TestTool(new AIToolDefinition { Name = "host", DisplayName = "Host", TenantBehavior = AITenantBehavior.HostScoped }), new TestTool(new AIToolDefinition { Name = "default-tenant-cross", DisplayName = "Default tenant cross", TenantBehavior = AITenantBehavior.CrossTenantDenied, TenantIds = [""] }) ]); var tools = await registry.ListAsync(new AIToolQuery { TenantId = "", ActorId = "user-1" }); Assert.Collection( tools.OrderBy(x => x.Name), tool => Assert.Equal("default-tenant-cross", tool.Name), tool => Assert.Equal("tenant", tool.Name)); } [Fact(DisplayName = "Tool registry allows cross-tenant denied tools for explicit tenant allowlists")] public async Task ToolRegistryAllowsCrossTenantDeniedToolsForExplicitTenantAllowlists() { var registry = CreateRegistry( [ new TestTool(new AIToolDefinition { Name = "tenant-cross", DisplayName = "Tenant cross", TenantBehavior = AITenantBehavior.CrossTenantDenied, TenantIds = ["tenant-1"] }) ]); var tools = await registry.ListAsync(new AIToolQuery { TenantId = "tenant-1", ActorId = "user-1" }); var tool = Assert.Single(tools); Assert.Equal("tenant-cross", tool.Name); } [Fact(DisplayName = "Tool registry excludes cross-tenant denied tools from host queries")] public async Task ToolRegistryExcludesCrossTenantDeniedToolsFromHostQueries() { var registry = CreateRegistry( [ new TestTool(new AIToolDefinition { Name = "host", DisplayName = "Host", TenantBehavior = AITenantBehavior.HostScoped }), new TestTool(new AIToolDefinition { Name = "cross", DisplayName = "Cross", TenantBehavior = AITenantBehavior.CrossTenantDenied }) ]); var tools = await registry.ListAsync(new AIToolQuery { ActorId = "user-1" }); var tool = Assert.Single(tools); Assert.Equal("host", tool.Name); } [Fact(DisplayName = "Tool registry exposes default tools without tenant context")] public async Task ToolRegistryExposesDefaultToolsWithoutTenantContext() { var registry = CreateRegistry( [ new TestTool(new AIToolDefinition { Name = "default", DisplayName = "Default" }) ]); var tools = await registry.ListAsync(new AIToolQuery { ActorId = "user-1" }); var tool = Assert.Single(tools); Assert.Equal("default", tool.Name); } [Fact(DisplayName = "Tool registry exposes default tools to tenant queries")] public async Task ToolRegistryExposesDefaultToolsToTenantQueries() { var registry = CreateRegistry( [ new TestTool(new AIToolDefinition { Name = "default", DisplayName = "Default" }) ]); var tools = await registry.ListAsync(new AIToolQuery { TenantId = "tenant-1", ActorId = "user-1" }); var tool = Assert.Single(tools); Assert.Equal("default", tool.Name); } [Fact(DisplayName = "Tool registry excludes incomplete tool definition names")] public async Task ToolRegistryExcludesIncompleteToolDefinitionNames() { var registry = CreateRegistry( [ new TestTool(new AIToolDefinition { DisplayName = "Incomplete" }), new TestTool(new AIToolDefinition { Name = "valid", DisplayName = "Valid" }) ]); var tools = await registry.ListAsync(new AIToolQuery { ActorId = "user-1" }); var tool = Assert.Single(tools); Assert.Equal("valid", tool.Name); } [Fact(DisplayName = "Tool registry uses cached concrete type lookup after listing tools")] public async Task ToolRegistryUsesCachedConcreteTypeLookupAfterListingTools() { CountingTool.ConstructionCount = 0; OtherCountingTool.ConstructionCount = 0; var services = new ServiceCollection(); services.AddTransient(); services.AddTransient(); var registry = CreateRegistry(services); await registry.ListAsync(new AIToolQuery { ActorId = "user-1" }); CountingTool.ConstructionCount = 0; OtherCountingTool.ConstructionCount = 0; using var tool = await registry.FindAsync("counting", new AIToolQuery { ActorId = "user-1" }); Assert.NotNull(tool); Assert.Equal(1, CountingTool.ConstructionCount); Assert.Equal(0, OtherCountingTool.ConstructionCount); } [Fact(DisplayName = "Tool registry disposes cached concrete tool instances")] public async Task ToolRegistryDisposesCachedConcreteToolInstances() { DisposableCachedTool.ConstructionCount = 0; DisposableCachedTool.DisposeCount = 0; var services = new ServiceCollection(); services.AddTransient(); var registry = CreateRegistry(services); await registry.ListAsync(new AIToolQuery { ActorId = "user-1" }); DisposableCachedTool.ConstructionCount = 0; DisposableCachedTool.DisposeCount = 0; var tool = await registry.FindAsync("disposable-cached", new AIToolQuery { ActorId = "user-1" }); tool?.Dispose(); Assert.NotNull(tool); Assert.Equal(1, DisposableCachedTool.ConstructionCount); Assert.Equal(1, DisposableCachedTool.DisposeCount); } [Fact(DisplayName = "Tool registry find applies tenant and actor filters")] public async Task ToolRegistryFindAppliesTenantAndActorFilters() { var registry = CreateRegistry( [ new TestTool(new AIToolDefinition { Name = "restricted", DisplayName = "Restricted", EnabledByDefault = true, TenantBehavior = AITenantBehavior.TenantScoped, TenantIds = ["tenant-1"], ActorIds = ["user-1"] }) ]); using var allowed = await registry.FindAsync("restricted", new AIToolQuery { TenantId = "tenant-1", ActorId = "user-1" }); var denied = await registry.FindAsync("restricted", new AIToolQuery { TenantId = "tenant-2", ActorId = "user-1" }); Assert.NotNull(allowed); Assert.Null(denied); } [Fact(DisplayName = "Tool registry enforces tool permission requirements")] public async Task ToolRegistryEnforcesToolPermissionRequirements() { var registry = CreateRegistry( [ new TestTool(new AIToolDefinition { Name = "restricted", DisplayName = "Restricted", EnabledByDefault = true, TenantBehavior = AITenantBehavior.TenantScoped, Permissions = ["workflows:write"] }) ]); var denied = await registry.FindAsync("restricted", new AIToolQuery { TenantId = "tenant-1", ActorId = "user-1" }); using var allowed = await registry.FindAsync("restricted", new AIToolQuery { TenantId = "tenant-1", ActorId = "user-1", UserPermissions = ["Workflows:Write"] }); Assert.Null(denied); Assert.NotNull(allowed); } [Fact(DisplayName = "Tool enablement requires explicit administrative enablement")] public void ToolEnablementRequiresExplicitAdministrativeEnablement() { var enablement = new AIToolEnablementService(); var definition = new AIToolDefinition { Name = "admin", DisplayName = "Admin", Mutability = AIToolMutability.Administrative, EnabledByDefault = true }; Assert.False(enablement.IsEnabled(definition)); enablement.Enable("admin"); Assert.False(enablement.IsEnabled(definition)); enablement.EnableAdministrative("admin"); Assert.True(enablement.IsEnabled(definition)); } [Fact(DisplayName = "Tool registry filters agent-scoped tools by agent")] public async Task ToolRegistryFiltersAgentScopedToolsByAgent() { var registry = CreateRegistry( [ new TestTool(new AIToolDefinition { Name = "agent-only", DisplayName = "Agent only", EnabledByDefault = true, TenantBehavior = AITenantBehavior.TenantScoped, AgentScopes = ["workflow-author"] }), new TestTool(new AIToolDefinition { Name = "agent-permission", DisplayName = "Agent permission", EnabledByDefault = true, TenantBehavior = AITenantBehavior.TenantScoped, AgentScopes = ["workflow-author"], Permissions = ["workflows:write"] }) ]); var deniedWithoutPermission = await registry.ListAsync(new AIToolQuery { Agent = "workflow-author", TenantId = "tenant-1", ActorId = "user-1" }); var deniedByAgent = await registry.ListAsync(new AIToolQuery { Agent = "workflow-editor", TenantId = "tenant-1", ActorId = "user-1", UserPermissions = ["workflows:write"] }); var allowedWithPermission = await registry.ListAsync(new AIToolQuery { Agent = "workflow-author", TenantId = "tenant-1", ActorId = "user-1", UserPermissions = ["workflows:write"] }); var tool = Assert.Single(deniedWithoutPermission); Assert.Equal("agent-only", tool.Name); Assert.Empty(deniedByAgent); Assert.Collection( allowedWithPermission.OrderBy(x => x.Name), agentOnly => Assert.Equal("agent-only", agentOnly.Name), agentPermission => Assert.Equal("agent-permission", agentPermission.Name)); } [Fact(DisplayName = "Tool registry honors tenant and actor allowlists")] public async Task ToolRegistryHonorsTenantAndActorAllowlists() { var registry = CreateRegistry( [ new TestTool(new AIToolDefinition { Name = "matching", DisplayName = "Matching", TenantBehavior = AITenantBehavior.TenantScoped, TenantIds = ["tenant-1"], ActorIds = ["user-1"] }), new TestTool(new AIToolDefinition { Name = "wrong-tenant", DisplayName = "Wrong tenant", TenantBehavior = AITenantBehavior.TenantScoped, TenantIds = ["tenant-2"] }), new TestTool(new AIToolDefinition { Name = "wrong-actor", DisplayName = "Wrong actor", ActorIds = ["user-2"] }) ]); var tools = await registry.ListAsync(new AIToolQuery { TenantId = "tenant-1", ActorId = "user-1" }); var tool = Assert.Single(tools); Assert.Equal("matching", tool.Name); } [Fact(DisplayName = "Tool registry combines tenant actor agent and permission filters")] public async Task ToolRegistryCombinesTenantActorAgentAndPermissionFilters() { var registry = CreateRegistry( [ new TestTool(new AIToolDefinition { Name = "matching", DisplayName = "Matching", TenantBehavior = AITenantBehavior.TenantScoped, TenantIds = ["tenant-1"], ActorIds = ["user-1"], AgentScopes = ["workflow-author"], Permissions = ["workflows:write"] }), new TestTool(new AIToolDefinition { Name = "wrong-agent", DisplayName = "Wrong agent", TenantBehavior = AITenantBehavior.TenantScoped, TenantIds = ["tenant-1"], ActorIds = ["user-1"], AgentScopes = ["workflow-editor"], Permissions = ["workflows:write"] }), new TestTool(new AIToolDefinition { Name = "wrong-permission", DisplayName = "Wrong permission", TenantBehavior = AITenantBehavior.TenantScoped, TenantIds = ["tenant-1"], ActorIds = ["user-1"], AgentScopes = ["workflow-author"], Permissions = ["deployments:write"] }) ]); var tools = await registry.ListAsync(new AIToolQuery { TenantId = "tenant-1", ActorId = "user-1", Agent = "workflow-author", UserPermissions = ["workflows:write"] }); var tool = Assert.Single(tools); Assert.Equal("matching", tool.Name); } [Fact(DisplayName = "Tool registry skips tools with throwing definitions")] public async Task ToolRegistrySkipsToolsWithThrowingDefinitions() { var tracker = new ScopeDisposalTracker(); var services = new ServiceCollection(); services.AddSingleton(tracker); services.AddScoped(); services.AddScoped(); services.AddScoped(_ => new TestTool(new AIToolDefinition { Name = "healthy", DisplayName = "Healthy" })); using var provider = services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true }); var registry = new AIToolRegistry(provider.GetRequiredService(), new AIToolEnablementService()); var missing = await registry.FindAsync("throwing", new AIToolQuery()); var tools = await registry.ListAsync(new AIToolQuery()); Assert.Null(missing); var tool = Assert.Single(tools); Assert.Equal("healthy", tool.Name); Assert.Equal(2, tracker.DisposeCount); } private static AIToolRegistry CreateRegistry(IReadOnlyCollection tools) { var services = new ServiceCollection(); foreach (var tool in tools) services.AddScoped(_ => tool); return CreateRegistry(services); } private static AIToolRegistry CreateRegistry(IServiceCollection services) { var provider = services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true }); return new AIToolRegistry(provider.GetRequiredService(), new AIToolEnablementService()); } private class TestTool(AIToolDefinition definition) : IAITool { public AIToolDefinition Definition { get; } = definition; public ValueTask ExecuteAsync(AIToolExecutionContext context, CancellationToken cancellationToken = default) => ValueTask.FromResult(new AIToolResult()); public void Dispose() { } } private class CountingTool : IAITool { public static int ConstructionCount { get; set; } public CountingTool() { ConstructionCount++; } public AIToolDefinition Definition { get; } = new() { Name = "counting", DisplayName = "Counting", EnabledByDefault = true }; public ValueTask ExecuteAsync(AIToolExecutionContext context, CancellationToken cancellationToken = default) => ValueTask.FromResult(new AIToolResult()); public void Dispose() { } } private class OtherCountingTool : IAITool { public static int ConstructionCount { get; set; } public OtherCountingTool() { ConstructionCount++; } public AIToolDefinition Definition { get; } = new() { Name = "other-counting", DisplayName = "Other counting", EnabledByDefault = true }; public ValueTask ExecuteAsync(AIToolExecutionContext context, CancellationToken cancellationToken = default) => ValueTask.FromResult(new AIToolResult()); public void Dispose() { } } private class DisposableCachedTool : IAITool { public static int ConstructionCount { get; set; } public static int DisposeCount { get; set; } public DisposableCachedTool() { ConstructionCount++; } public AIToolDefinition Definition { get; } = new() { Name = "disposable-cached", DisplayName = "Disposable cached", EnabledByDefault = true }; public ValueTask ExecuteAsync(AIToolExecutionContext context, CancellationToken cancellationToken = default) => ValueTask.FromResult(new AIToolResult()); public void Dispose() { DisposeCount++; } } private class ThrowingDefinitionTool(ScopedDependency dependency) : IAITool { private readonly ScopedDependency _dependency = dependency; public AIToolDefinition Definition { get { _ = _dependency; throw new InvalidOperationException("Definition unavailable."); } } public ValueTask ExecuteAsync(AIToolExecutionContext context, CancellationToken cancellationToken = default) => ValueTask.FromResult(new AIToolResult()); public void Dispose() { } } private class ScopedDependency(ScopeDisposalTracker tracker) : IDisposable { public void Dispose() { tracker.DisposeCount++; } } private class ScopeDisposalTracker { public int DisposeCount { get; set; } } }