From 937df080debab9abfb9ed73eb3de3f1ffb95e79c Mon Sep 17 00:00:00 2001 From: Adam Frisby Date: Mon, 14 Sep 2026 04:24:17 +0000 Subject: [PATCH] Fix /quota to resolve probes via AgentQuotaProbeCatalog like dispatch The report used a Kind-keyed dictionary, so members served by a probe whose Kind differs (e.g. zen/go copilot via the opencode-go probe) were silently omitted. Resolve each member with ResolveSubscriptionProbe, surface conflicts as unknown snapshots, and keep omitting genuinely unclaimed members. CodeyBox-Prompt-Revision: 1 Co-Authored-By: CodeyBox --- src/CodeyBox.Api/Program.cs | 33 ++-- .../AgentQuotaProbeCatalog.cs | 6 +- tests/CodeyBox.Tests/QuotaEndpointTests.cs | 150 ++++++++++++++++++ 3 files changed, 175 insertions(+), 14 deletions(-) diff --git a/src/CodeyBox.Api/Program.cs b/src/CodeyBox.Api/Program.cs index e99492111..9863691c8 100644 --- a/src/CodeyBox.Api/Program.cs +++ b/src/CodeyBox.Api/Program.cs @@ -4433,20 +4433,31 @@ await ctx.Response.WriteAsync( var pausedByAgent = pausedStates .Where(s => s.AgentInstanceId is null) .ToDictionary(s => s.Agent, s => s); - var probeByKind = probes - .Where(p => p is not PayPerApiQuotaProbe and not NullQuotaProbe) - .ToDictionary(p => p.Kind); + var quotaLog = loggerFactory.CreateLogger("Quota"); + var subscriptionProbes = AgentQuotaProbeCatalog.BuildSubscriptionProbes(probes); + var fallbackKinds = subscriptionProbes.Select(p => p.Kind).Distinct().ToList(); var representedProbeKeys = new HashSet<(AgentKind Agent, string? ModelId)>(); var snapshots = new List(); var kindAggregateCounts = new Dictionary(StringComparer.OrdinalIgnoreCase); async Task AddSnapshotAsync(AgentMembership member, string? classId, string? classDisplayName) { - if (!probeByKind.TryGetValue(member.Agent, out var probe)) + var resolution = AgentQuotaProbeCatalog.ResolveSubscriptionProbe(subscriptionProbes, member, quotaLog); + AgentQuotaSnapshot snapshot; + if (resolution.Conflict is not null) + { + representedProbeKeys.Add((member.Agent, member.ModelId)); + snapshot = AgentQuotaProbeCatalog.ConflictUnknownSnapshot(resolution.Conflict); + } + else if (resolution.Probe is null) + { return; - - representedProbeKeys.Add((member.Agent, member.ModelId)); - var snapshot = await probe.GetAvailabilityAsync(member, ct); + } + else + { + representedProbeKeys.Add((member.Agent, member.ModelId)); + snapshot = await resolution.Probe.GetAvailabilityAsync(member, ct); + } var poolName = QuotaPoolResolver.NormalizePoolName(member.Pool); string? poolKind = null; if (poolName is not null @@ -4516,7 +4527,7 @@ bool WouldAllow(AgentMembership gateMember, bool hasRecentFailure) => if (paused) return false; var modelMember = member with { ModelId = modelId }; var modelHasRecentFailure = recentFailuresForProbe.Any(f => - f.Agent == probe.Kind && + f.Agent == member.Agent && string.Equals(f.ModelId, modelId, StringComparison.OrdinalIgnoreCase)); return WouldAllow(modelMember, modelHasRecentFailure); }, @@ -4539,14 +4550,14 @@ bool WouldAllow(AgentMembership gateMember, bool hasRecentFailure) => } } - foreach (var probe in probeByKind.Values) + foreach (var kind in fallbackKinds) { - if (representedProbeKeys.Any(k => k.Agent == probe.Kind && k.ModelId is null)) + if (representedProbeKeys.Any(k => k.Agent == kind && k.ModelId is null)) continue; await AddSnapshotAsync(new AgentMembership { - Agent = probe.Kind, + Agent = kind, Billing = AgentBilling.Subscription, QualityScore = 100, }, classId: null, classDisplayName: null); diff --git a/src/CodeyBox.Orchestrator/AgentQuotaProbeCatalog.cs b/src/CodeyBox.Orchestrator/AgentQuotaProbeCatalog.cs index 2220ac343..45ae29341 100644 --- a/src/CodeyBox.Orchestrator/AgentQuotaProbeCatalog.cs +++ b/src/CodeyBox.Orchestrator/AgentQuotaProbeCatalog.cs @@ -10,18 +10,18 @@ namespace CodeyBox.Orchestrator; /// must fail closed for that key); both are null when no probe claims the member /// (callers fall back to the NullQuotaProbe unknown path, as before). /// -internal sealed record QuotaProbeResolution(IAgentQuotaProbe? Probe, QuotaProbeConflict? Conflict); +public sealed record QuotaProbeResolution(IAgentQuotaProbe? Probe, QuotaProbeConflict? Conflict); /// /// Two or more equally specific probes claimed the same member. Carries the /// contested key and the tied probe identities for the Error log and the /// fail-closed unknown snapshot. /// -internal sealed record QuotaProbeConflict( +public sealed record QuotaProbeConflict( AgentQuotaMemberKey Key, IReadOnlyList ProbeNames); -internal static class AgentQuotaProbeCatalog +public static class AgentQuotaProbeCatalog { /// /// Sentinel model id used only to rank claim specificity: no real routing diff --git a/tests/CodeyBox.Tests/QuotaEndpointTests.cs b/tests/CodeyBox.Tests/QuotaEndpointTests.cs index ef717c3be..c46caa846 100644 --- a/tests/CodeyBox.Tests/QuotaEndpointTests.cs +++ b/tests/CodeyBox.Tests/QuotaEndpointTests.cs @@ -301,6 +301,156 @@ public async Task GetQuota_IncludesConfiguredInstancesSeparatelyAndKindAggregate Assert.Equal(2, aggregate.GetProperty("instances").GetInt32()); } + private sealed class CrossKindProbe : IAgentQuotaProbe + { + private readonly AgentQuotaSnapshot _snapshot; + private readonly AgentKind _claimed; + + public CrossKindProbe(AgentKind kind, AgentKind claimed, AgentQuotaSnapshot snapshot) + { + Kind = kind; + _claimed = claimed; + _snapshot = snapshot; + } + + public AgentKind Kind { get; } + + public bool Handles(AgentQuotaMemberKey key) => key.Agent == _claimed; + + public Task GetAvailabilityAsync(AgentMembership member, CancellationToken ct) + => Task.FromResult(_snapshot); + } + + [Fact] + public async Task GetQuota_ReportsMemberServedByCrossKindProbe() + { + using var factory = new WorkItemApiFactory(); + var client = factory.WithWebHostBuilder(builder => + { + builder.ConfigureAppConfiguration((_, cfg) => + { + cfg.AddInMemoryCollection(new Dictionary + { + ["CodeyBox:AgentClasses:0:Id"] = "go", + ["CodeyBox:AgentClasses:0:DisplayName"] = "Go", + ["CodeyBox:AgentClasses:0:Members:0:Agent"] = "copilot", + ["CodeyBox:AgentClasses:0:Members:0:Billing"] = "Subscription", + ["CodeyBox:AgentClasses:0:Members:0:QualityScore"] = "100", + }); + }); + builder.ConfigureTestServices(services => + { + services.RemoveAll(); + services.AddSingleton(new CrossKindProbe( + AgentKind.Opencode, + AgentKind.Copilot, + new AgentQuotaSnapshot { AvailablePct = 57 })); + }); + }).CreateClient(); + + var response = await client.GetAsync("/quota"); + response.EnsureSuccessStatusCode(); + using var doc = JsonDocument.Parse(await response.Content.ReadAsStringAsync()); + + var root = doc.RootElement; + var copilotRows = root.GetProperty("probes") + .EnumerateArray() + .Where(p => string.Equals(p.GetProperty("agent").GetString(), "copilot", StringComparison.OrdinalIgnoreCase)) + .ToList(); + var row = Assert.Single(copilotRows); + Assert.Equal(57, row.GetProperty("latestSnapshot").GetProperty("availablePct").GetDouble()); + + var aggregate = Assert.Single(root.GetProperty("kindAggregates").EnumerateArray(), a => + string.Equals(a.GetProperty("agent").GetString(), "copilot", StringComparison.OrdinalIgnoreCase)); + Assert.Equal(1, aggregate.GetProperty("instances").GetInt32()); + } + + [Fact] + public async Task GetQuota_OmitsMemberNoProbeClaims() + { + using var factory = new WorkItemApiFactory(); + var client = factory.WithWebHostBuilder(builder => + { + builder.ConfigureAppConfiguration((_, cfg) => + { + cfg.AddInMemoryCollection(new Dictionary + { + ["CodeyBox:AgentClasses:0:Id"] = "mix", + ["CodeyBox:AgentClasses:0:DisplayName"] = "Mix", + ["CodeyBox:AgentClasses:0:Members:0:Agent"] = "copilot", + ["CodeyBox:AgentClasses:0:Members:0:Billing"] = "Subscription", + ["CodeyBox:AgentClasses:0:Members:0:QualityScore"] = "100", + }); + }); + builder.ConfigureTestServices(services => + { + services.RemoveAll(); + services.AddSingleton(new FakeProbe(AgentKind.Claude, 60)); + }); + }).CreateClient(); + + var response = await client.GetAsync("/quota"); + var debugBody = await response.Content.ReadAsStringAsync(); + Assert.True(response.IsSuccessStatusCode, debugBody); + using var doc = JsonDocument.Parse(debugBody); + + var root = doc.RootElement; + Assert.DoesNotContain( + root.GetProperty("probes").EnumerateArray(), + p => string.Equals(p.GetProperty("agent").GetString(), "copilot", StringComparison.OrdinalIgnoreCase)); + Assert.DoesNotContain( + root.GetProperty("kindAggregates").EnumerateArray(), + a => string.Equals(a.GetProperty("agent").GetString(), "copilot", StringComparison.OrdinalIgnoreCase)); + } + + [Fact] + public async Task GetQuota_ConflictingProbesReportUnknownInsteadOfDropping() + { + using var factory = new WorkItemApiFactory(); + var client = factory.WithWebHostBuilder(builder => + { + builder.ConfigureAppConfiguration((_, cfg) => + { + cfg.AddInMemoryCollection(new Dictionary + { + ["CodeyBox:AgentClasses:0:Id"] = "go", + ["CodeyBox:AgentClasses:0:DisplayName"] = "Go", + ["CodeyBox:AgentClasses:0:Members:0:Agent"] = "copilot", + ["CodeyBox:AgentClasses:0:Members:0:Billing"] = "Subscription", + ["CodeyBox:AgentClasses:0:Members:0:QualityScore"] = "100", + }); + }); + builder.ConfigureTestServices(services => + { + services.RemoveAll(); + services.AddSingleton(new CrossKindProbe( + AgentKind.Opencode, + AgentKind.Copilot, + new AgentQuotaSnapshot { AvailablePct = 57 })); + services.AddSingleton(new CrossKindProbe( + AgentKind.Claude, + AgentKind.Copilot, + new AgentQuotaSnapshot { AvailablePct = 80 })); + }); + }).CreateClient(); + + var response = await client.GetAsync("/quota"); + response.EnsureSuccessStatusCode(); + using var doc = JsonDocument.Parse(await response.Content.ReadAsStringAsync()); + + var root = doc.RootElement; + var row = Assert.Single( + root.GetProperty("probes").EnumerateArray(), + p => string.Equals(p.GetProperty("agent").GetString(), "copilot", StringComparison.OrdinalIgnoreCase)); + var latest = row.GetProperty("latestSnapshot"); + Assert.Equal(-1, latest.GetProperty("availablePct").GetDouble()); + Assert.Contains("conflicting quota probes", latest.GetProperty("notes").GetString(), StringComparison.OrdinalIgnoreCase); + + var aggregate = Assert.Single(root.GetProperty("kindAggregates").EnumerateArray(), a => + string.Equals(a.GetProperty("agent").GetString(), "copilot", StringComparison.OrdinalIgnoreCase)); + Assert.Equal(1, aggregate.GetProperty("instances").GetInt32()); + } + private sealed class FakeBudgetProvider : IAgentBudgetProvider { public Task GetBudgetSnapshotAsync(AgentKind agent, string? modelId, CancellationToken ct = default)