Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 22 additions & 11 deletions src/CodeyBox.Api/Program.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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<object>();
var kindAggregateCounts = new Dictionary<string, int>(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
Expand Down Expand Up @@ -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);
},
Expand All @@ -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);
Expand Down
6 changes: 3 additions & 3 deletions src/CodeyBox.Orchestrator/AgentQuotaProbeCatalog.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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 <c>NullQuotaProbe</c> unknown path, as before).
/// </summary>
internal sealed record QuotaProbeResolution(IAgentQuotaProbe? Probe, QuotaProbeConflict? Conflict);
public sealed record QuotaProbeResolution(IAgentQuotaProbe? Probe, QuotaProbeConflict? Conflict);

/// <summary>
/// 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.
/// </summary>
internal sealed record QuotaProbeConflict(
public sealed record QuotaProbeConflict(
AgentQuotaMemberKey Key,
IReadOnlyList<string> ProbeNames);

internal static class AgentQuotaProbeCatalog
public static class AgentQuotaProbeCatalog
{
/// <summary>
/// Sentinel model id used only to rank claim specificity: no real routing
Expand Down
150 changes: 150 additions & 0 deletions tests/CodeyBox.Tests/QuotaEndpointTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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<AgentQuotaSnapshot> 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<string, string?>
{
["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<IAgentQuotaProbe>();
services.AddSingleton<IAgentQuotaProbe>(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<string, string?>
{
["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<IAgentQuotaProbe>();
services.AddSingleton<IAgentQuotaProbe>(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<string, string?>
{
["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<IAgentQuotaProbe>();
services.AddSingleton<IAgentQuotaProbe>(new CrossKindProbe(
AgentKind.Opencode,
AgentKind.Copilot,
new AgentQuotaSnapshot { AvailablePct = 57 }));
services.AddSingleton<IAgentQuotaProbe>(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<AgentQuotaSnapshot?> GetBudgetSnapshotAsync(AgentKind agent, string? modelId, CancellationToken ct = default)
Expand Down
Loading