| | | 1 | | using NexusLabs.Needlr.AgentFramework.Context; |
| | | 2 | | |
| | | 3 | | namespace NexusLabs.Needlr.AgentFramework.Providers; |
| | | 4 | | |
| | | 5 | | /// <summary> |
| | | 6 | | /// Default <see cref="ITieredProviderSelector{TQuery, TResult}"/> that iterates providers |
| | | 7 | | /// in ascending <see cref="ITieredProvider{TQuery, TResult}.Priority"/> order, gated by |
| | | 8 | | /// an <see cref="IQuotaGate"/>. On <see cref="ProviderUnavailableException"/> the selector |
| | | 9 | | /// falls through to the next provider. |
| | | 10 | | /// </summary> |
| | | 11 | | /// <remarks> |
| | | 12 | | /// <para> |
| | | 13 | | /// The quota partition key is resolved from the ambient |
| | | 14 | | /// <see cref="IAgentExecutionContextAccessor"/> using a <see cref="QuotaPartitionSelector"/>. |
| | | 15 | | /// By default, <see cref="IAgentExecutionContext.UserId"/> is used as the partition. |
| | | 16 | | /// Consumers that need a different partitioning strategy (e.g., tenant ID, API key) |
| | | 17 | | /// can provide a custom <see cref="QuotaPartitionSelector"/> via the constructor. |
| | | 18 | | /// </para> |
| | | 19 | | /// <para> |
| | | 20 | | /// When no execution context is active (e.g., during integration tests that don't |
| | | 21 | | /// establish a scope), the partition is <see langword="null"/> and quota is global. |
| | | 22 | | /// </para> |
| | | 23 | | /// </remarks> |
| | | 24 | | [DoNotAutoRegister] |
| | | 25 | | public sealed class TieredProviderSelector<TQuery, TResult> : ITieredProviderSelector<TQuery, TResult> |
| | | 26 | | { |
| | | 27 | | private readonly IReadOnlyList<ITieredProvider<TQuery, TResult>> _providers; |
| | | 28 | | private readonly IQuotaGate _quotaGate; |
| | | 29 | | private readonly IAgentExecutionContextAccessor _contextAccessor; |
| | | 30 | | private readonly QuotaPartitionSelector _partitionSelector; |
| | | 31 | | |
| | | 32 | | /// <summary> |
| | | 33 | | /// The default partition selector: uses <see cref="IAgentExecutionContext.UserId"/> |
| | | 34 | | /// from the ambient context. |
| | | 35 | | /// </summary> |
| | 1 | 36 | | public static readonly QuotaPartitionSelector DefaultPartitionSelector = |
| | 12 | 37 | | context => context?.UserId; |
| | | 38 | | |
| | | 39 | | /// <param name="providers">All registered providers (filtering and ordering is handled internally).</param> |
| | | 40 | | /// <param name="quotaGate">Quota gate for reservation/release. Use <see cref="AlwaysGrantQuotaGate"/> for no-op.</p |
| | | 41 | | /// <param name="contextAccessor">Accessor for ambient execution context (provides partition identity).</param> |
| | | 42 | | /// <param name="partitionSelector"> |
| | | 43 | | /// Custom partition selector. Defaults to <see cref="DefaultPartitionSelector"/> |
| | | 44 | | /// (<see cref="IAgentExecutionContext.UserId"/>). |
| | | 45 | | /// </param> |
| | 16 | 46 | | public TieredProviderSelector( |
| | 16 | 47 | | IEnumerable<ITieredProvider<TQuery, TResult>> providers, |
| | 16 | 48 | | IQuotaGate quotaGate, |
| | 16 | 49 | | IAgentExecutionContextAccessor contextAccessor, |
| | 16 | 50 | | QuotaPartitionSelector? partitionSelector = null) |
| | | 51 | | { |
| | 16 | 52 | | ArgumentNullException.ThrowIfNull(providers); |
| | 15 | 53 | | ArgumentNullException.ThrowIfNull(quotaGate); |
| | 14 | 54 | | ArgumentNullException.ThrowIfNull(contextAccessor); |
| | | 55 | | |
| | 13 | 56 | | _providers = providers |
| | 20 | 57 | | .Where(p => p.IsEnabled) |
| | 13 | 58 | | .OrderBy(p => p.Priority) |
| | 13 | 59 | | .ThenBy(p => p.Name, StringComparer.OrdinalIgnoreCase) |
| | 13 | 60 | | .ToList(); |
| | | 61 | | |
| | 13 | 62 | | _quotaGate = quotaGate; |
| | 13 | 63 | | _contextAccessor = contextAccessor; |
| | 13 | 64 | | _partitionSelector = partitionSelector ?? DefaultPartitionSelector; |
| | 13 | 65 | | } |
| | | 66 | | |
| | | 67 | | /// <inheritdoc /> |
| | | 68 | | public async Task<TResult> ExecuteAsync(TQuery query, CancellationToken cancellationToken) |
| | | 69 | | { |
| | 13 | 70 | | if (_providers.Count == 0) |
| | 1 | 71 | | throw new InvalidOperationException("No enabled providers are registered."); |
| | | 72 | | |
| | 12 | 73 | | var partition = _partitionSelector(_contextAccessor.Current); |
| | 12 | 74 | | var attempts = new List<string>(); |
| | | 75 | | |
| | 45 | 76 | | foreach (var provider in _providers) |
| | | 77 | | { |
| | 16 | 78 | | if (!await _quotaGate.TryReserveAsync(provider.Name, partition, cancellationToken) |
| | 16 | 79 | | .ConfigureAwait(false)) |
| | | 80 | | { |
| | 1 | 81 | | attempts.Add($"{provider.Name}: quota denied"); |
| | 1 | 82 | | continue; |
| | | 83 | | } |
| | | 84 | | |
| | | 85 | | try |
| | | 86 | | { |
| | 15 | 87 | | var result = await provider.ExecuteAsync(query, cancellationToken) |
| | 15 | 88 | | .ConfigureAwait(false); |
| | | 89 | | |
| | 11 | 90 | | await _quotaGate.ReleaseAsync(provider.Name, partition, succeeded: true, cancellationToken) |
| | 11 | 91 | | .ConfigureAwait(false); |
| | | 92 | | |
| | 11 | 93 | | return result; |
| | | 94 | | } |
| | | 95 | | catch (ProviderUnavailableException ex) |
| | | 96 | | { |
| | 4 | 97 | | attempts.Add($"{provider.Name}: {ex.Message}"); |
| | | 98 | | |
| | 4 | 99 | | await _quotaGate.ReleaseAsync(provider.Name, partition, succeeded: false, cancellationToken) |
| | 4 | 100 | | .ConfigureAwait(false); |
| | | 101 | | } |
| | 4 | 102 | | } |
| | | 103 | | |
| | 1 | 104 | | throw new InvalidOperationException( |
| | 1 | 105 | | $"All providers failed. Attempts: [{string.Join("; ", attempts)}]"); |
| | 11 | 106 | | } |
| | | 107 | | } |