Budgetende, Stufe 2: laufende Laeufe reservieren ihr Budget (#716)
Der offene Rest aus #682: Jeder Lauf pruefte das Budget einzeln gegen den bisher erfassten Verbrauch. Laufende Laeufe waren unsichtbar, gleichzeitig startende passierten die Pruefung gemeinsam und ueberzogen gemeinsam. - BudgetGuard.ReserveAsync prueft und reserviert unter einer Sperre. Die Reservierung ist der Mittelwert der letzten 20 Laeufe des Agenten (Laeufe ohne Verbrauch zaehlen nicht mit) und gilt, bis der Verbrauch des Laufs erfasst ist. - Gegen die Grenze zaehlt: verbraucht + in laufenden Laeufen reserviert. Die eigene Schaetzung eines Laufs zaehlt bewusst nicht mit; ein einzelner Lauf wird weiter nur abgelehnt, wenn das Budget schon verbraucht ist. Sonst fiele eine kleine Aufgabe wegen eines grossen Durchschnitts aus. - Gilt fuer Kosten- und Tokengrenze, je Agent und je Instanz. - Die Ablehnungsmeldung nennt den reservierten Anteil, sonst wirkte "0,60 von 1,00 USD verbraucht" als Ablehnungsgrund falsch. - IUsageRepository.GetRecentForAgentAsync, mit leerer Vorgabe im Interface, damit bestehende Attrappen unberuehrt bleiben. Tests: BudgetReservationTests (8 neu), darunter zehn gleichzeitige Starts, von denen genau vier passieren. 1.246 gruen. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5.5
parent
5f3a85595b
commit
8a83a67fb4
@@ -125,6 +125,29 @@ public sealed class SqliteUsageRepository : IUsageRepository
|
||||
return results;
|
||||
}
|
||||
|
||||
public async Task<IReadOnlyList<RunUsage>> GetRecentForAgentAsync(string agentId, int limit, CancellationToken ct)
|
||||
{
|
||||
await using var conn = await _storage.OpenConnectionAsync(ct);
|
||||
using var cmd = conn.CreateCommand();
|
||||
cmd.CommandText = """
|
||||
SELECT Id, AgentId, Model, PromptTokens, CompletionTokens, CachedTokens,
|
||||
CostUsd, CostIsKnown, Status, StepCount, DurationMs, OccurredAt
|
||||
FROM RunUsage
|
||||
WHERE AgentId = @agent COLLATE NOCASE
|
||||
ORDER BY OccurredAt DESC
|
||||
LIMIT @limit
|
||||
""";
|
||||
cmd.Parameters.AddWithValue("@agent", agentId);
|
||||
cmd.Parameters.AddWithValue("@limit", Math.Clamp(limit, 1, 1000));
|
||||
|
||||
var results = new List<RunUsage>();
|
||||
await using var reader = await cmd.ExecuteReaderAsync(ct);
|
||||
while (await reader.ReadAsync(ct))
|
||||
results.Add(Read(reader));
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
public Task<int> PurgeBeforeAsync(DateOnly date, CancellationToken ct)
|
||||
=> _storage.WriteAsync(async conn =>
|
||||
{
|
||||
|
||||
@@ -53,6 +53,14 @@ public interface IUsageRepository
|
||||
|
||||
Task<IReadOnlyList<RunUsage>> GetRecentAsync(int limit, CancellationToken ct);
|
||||
|
||||
/// <summary>
|
||||
/// #716: The most recent runs of one agent — the basis for estimating what a run of
|
||||
/// it will cost. Without an implementation there is no estimate, and a reservation
|
||||
/// counts nothing.
|
||||
/// </summary>
|
||||
Task<IReadOnlyList<RunUsage>> GetRecentForAgentAsync(string agentId, int limit, CancellationToken ct)
|
||||
=> Task.FromResult<IReadOnlyList<RunUsage>>([]);
|
||||
|
||||
/// <summary>Entfernt Einträge älter als das angegebene Datum. Gibt die Anzahl zurück.</summary>
|
||||
Task<int> PurgeBeforeAsync(DateOnly date, CancellationToken ct);
|
||||
}
|
||||
|
||||
@@ -13,6 +13,13 @@ public enum BudgetLimit
|
||||
InstanceTokens
|
||||
}
|
||||
|
||||
/// <summary>#716: What a run is expected to use — the agent's average over recent runs.</summary>
|
||||
public readonly record struct RunEstimate(decimal CostUsd, long Tokens)
|
||||
{
|
||||
public static RunEstimate operator +(RunEstimate a, RunEstimate b) => new(a.CostUsd + b.CostUsd, a.Tokens + b.Tokens);
|
||||
public static RunEstimate operator -(RunEstimate a, RunEstimate b) => new(a.CostUsd - b.CostUsd, a.Tokens - b.Tokens);
|
||||
}
|
||||
|
||||
public sealed record BudgetStatus(bool IsAllowed, BudgetLimit Limit, string? Reason)
|
||||
{
|
||||
public static BudgetStatus Allowed { get; } = new(true, BudgetLimit.None, null);
|
||||
@@ -34,6 +41,19 @@ public sealed class BudgetGuard
|
||||
private readonly IUsageRepository _usage;
|
||||
private readonly Func<DateTime> _now;
|
||||
|
||||
/// <summary>How many recent runs of an agent make up its per-run estimate (#716).</summary>
|
||||
public const int EstimateSample = 20;
|
||||
|
||||
/// <summary>
|
||||
/// #716: Check and reserve happen under this lock. Without it, two runs starting at
|
||||
/// the same moment would both read the same "used" figure and both pass.
|
||||
/// </summary>
|
||||
private readonly SemaphoreSlim _reserveGate = new(1, 1);
|
||||
|
||||
private readonly Lock _reservedLock = new();
|
||||
private readonly Dictionary<string, RunEstimate> _reservedByAgent = new(StringComparer.OrdinalIgnoreCase);
|
||||
private RunEstimate _reservedInstance;
|
||||
|
||||
public BudgetGuard(IUsageRepository usage, Func<DateTime>? now = null)
|
||||
{
|
||||
_usage = usage;
|
||||
@@ -43,8 +63,112 @@ public sealed class BudgetGuard
|
||||
/// <summary>
|
||||
/// Prüft, ob der Agent heute noch laufen darf.
|
||||
/// </summary>
|
||||
public async Task<BudgetStatus> CheckAsync(
|
||||
public Task<BudgetStatus> CheckAsync(
|
||||
AgentConfig agent, InstanceBudget instanceBudget, CancellationToken ct)
|
||||
=> CheckCoreAsync(agent, instanceBudget, default, default, ct);
|
||||
|
||||
/// <summary>
|
||||
/// #716: Checks like <see cref="CheckAsync"/> and, if the run may start, reserves its
|
||||
/// estimated usage until the returned handle is disposed — i.e. until the run's usage
|
||||
/// is recorded.
|
||||
///
|
||||
/// <para><b>What changes.</b> Before, each run was checked on its own against what was
|
||||
/// already recorded. Runs in flight were invisible, so several starting together all
|
||||
/// passed and exceeded the budget together (#682 capped how many; this closes it).
|
||||
/// Now the runs in flight count with their estimate.</para>
|
||||
///
|
||||
/// <para><b>What does not change.</b> A single run is still refused only when the
|
||||
/// budget is already used up — its own estimate is not added. Otherwise a small task
|
||||
/// would be refused because the agent's average run is large.</para>
|
||||
/// </summary>
|
||||
public async Task<(BudgetStatus Status, IDisposable? Reservation)> ReserveAsync(
|
||||
AgentConfig agent, InstanceBudget instanceBudget, CancellationToken ct)
|
||||
{
|
||||
if (!agent.Budget.HasAnyLimit && !instanceBudget.HasAnyLimit)
|
||||
return (BudgetStatus.Allowed, null);
|
||||
|
||||
var estimate = await EstimateRunAsync(agent.AgentId, ct);
|
||||
|
||||
await _reserveGate.WaitAsync(ct);
|
||||
try
|
||||
{
|
||||
RunEstimate agentReserved, instanceReserved;
|
||||
lock (_reservedLock)
|
||||
{
|
||||
agentReserved = _reservedByAgent.GetValueOrDefault(agent.AgentId);
|
||||
instanceReserved = _reservedInstance;
|
||||
}
|
||||
|
||||
var status = await CheckCoreAsync(agent, instanceBudget, agentReserved, instanceReserved, ct);
|
||||
if (!status.IsAllowed)
|
||||
return (status, null);
|
||||
|
||||
lock (_reservedLock)
|
||||
{
|
||||
_reservedByAgent[agent.AgentId] = _reservedByAgent.GetValueOrDefault(agent.AgentId) + estimate;
|
||||
_reservedInstance += estimate;
|
||||
}
|
||||
|
||||
return (status, new Reservation(this, agent.AgentId, estimate));
|
||||
}
|
||||
finally
|
||||
{
|
||||
_reserveGate.Release();
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>What is currently reserved by runs in flight — for the agent, or with an
|
||||
/// empty id for the whole instance.</summary>
|
||||
public RunEstimate Reserved(string agentId)
|
||||
{
|
||||
lock (_reservedLock)
|
||||
return agentId.Length == 0 ? _reservedInstance : _reservedByAgent.GetValueOrDefault(agentId);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// The average of the agent's recent runs. Runs without usage (refusals, aborted
|
||||
/// before the first call) do not count — they would drag the estimate towards zero.
|
||||
/// </summary>
|
||||
public async Task<RunEstimate> EstimateRunAsync(string agentId, CancellationToken ct)
|
||||
{
|
||||
var recent = await _usage.GetRecentForAgentAsync(agentId, EstimateSample, ct);
|
||||
var runs = recent.Where(r => r.TotalTokens > 0).ToList();
|
||||
if (runs.Count == 0)
|
||||
return default;
|
||||
|
||||
return new RunEstimate(
|
||||
runs.Average(r => r.CostUsd),
|
||||
(long)Math.Ceiling(runs.Average(r => (double)r.TotalTokens)));
|
||||
}
|
||||
|
||||
private void Release(string agentId, RunEstimate estimate)
|
||||
{
|
||||
lock (_reservedLock)
|
||||
{
|
||||
var left = _reservedByAgent.GetValueOrDefault(agentId) - estimate;
|
||||
if (left.CostUsd <= 0 && left.Tokens <= 0)
|
||||
_reservedByAgent.Remove(agentId);
|
||||
else
|
||||
_reservedByAgent[agentId] = left;
|
||||
|
||||
_reservedInstance -= estimate;
|
||||
}
|
||||
}
|
||||
|
||||
private sealed class Reservation(BudgetGuard guard, string agentId, RunEstimate estimate) : IDisposable
|
||||
{
|
||||
private int _disposed;
|
||||
|
||||
public void Dispose()
|
||||
{
|
||||
if (Interlocked.Exchange(ref _disposed, 1) == 0)
|
||||
guard.Release(agentId, estimate);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task<BudgetStatus> CheckCoreAsync(
|
||||
AgentConfig agent, InstanceBudget instanceBudget,
|
||||
RunEstimate agentReserved, RunEstimate instanceReserved, CancellationToken ct)
|
||||
{
|
||||
var today = DateOnly.FromDateTime(_now());
|
||||
|
||||
@@ -54,19 +178,21 @@ public sealed class BudgetGuard
|
||||
{
|
||||
var used = await _usage.GetDailyAsync(today, agent.AgentId, ct);
|
||||
|
||||
if (Exceeds(agentBudget.DailyCostUsd, used.CostUsd))
|
||||
if (Exceeds(agentBudget.DailyCostUsd, used.CostUsd + agentReserved.CostUsd))
|
||||
{
|
||||
return new BudgetStatus(false, BudgetLimit.AgentCost,
|
||||
$"Tagesbudget des Agenten erschöpft: {used.CostUsd:F4} von maximal " +
|
||||
$"{agentBudget.DailyCostUsd:F2} USD verbraucht." +
|
||||
ReservedHint(agentReserved.CostUsd) +
|
||||
UnpricedHint(used));
|
||||
}
|
||||
|
||||
if (Exceeds(agentBudget.DailyTokens, used.TotalTokens))
|
||||
if (Exceeds(agentBudget.DailyTokens, used.TotalTokens + agentReserved.Tokens))
|
||||
{
|
||||
return new BudgetStatus(false, BudgetLimit.AgentTokens,
|
||||
$"Token-Tageslimit des Agenten erschöpft: {used.TotalTokens:N0} von maximal " +
|
||||
$"{agentBudget.DailyTokens:N0}.");
|
||||
$"{agentBudget.DailyTokens:N0}." +
|
||||
ReservedTokensHint(agentReserved.Tokens));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -75,19 +201,21 @@ public sealed class BudgetGuard
|
||||
{
|
||||
var used = await _usage.GetDailyAsync(today, agentId: "", ct);
|
||||
|
||||
if (Exceeds(instanceBudget.DailyCostUsd, used.CostUsd))
|
||||
if (Exceeds(instanceBudget.DailyCostUsd, used.CostUsd + instanceReserved.CostUsd))
|
||||
{
|
||||
return new BudgetStatus(false, BudgetLimit.InstanceCost,
|
||||
$"Tagesbudget der Instanz erschöpft: {used.CostUsd:F4} von maximal " +
|
||||
$"{instanceBudget.DailyCostUsd:F2} USD verbraucht." +
|
||||
ReservedHint(instanceReserved.CostUsd) +
|
||||
UnpricedHint(used));
|
||||
}
|
||||
|
||||
if (Exceeds(instanceBudget.DailyTokens, used.TotalTokens))
|
||||
if (Exceeds(instanceBudget.DailyTokens, used.TotalTokens + instanceReserved.Tokens))
|
||||
{
|
||||
return new BudgetStatus(false, BudgetLimit.InstanceTokens,
|
||||
$"Token-Tageslimit der Instanz erschöpft: {used.TotalTokens:N0} von maximal " +
|
||||
$"{instanceBudget.DailyTokens:N0}.");
|
||||
$"{instanceBudget.DailyTokens:N0}." +
|
||||
ReservedTokensHint(instanceReserved.Tokens));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -109,6 +237,13 @@ public sealed class BudgetGuard
|
||||
|
||||
private static bool Exceeds(long limit, long used) => limit > 0 && used >= limit;
|
||||
|
||||
/// <summary>#716: Without this, a refusal at "3.20 of 5.00 USD used" would look wrong.</summary>
|
||||
private static string ReservedHint(decimal reserved)
|
||||
=> reserved > 0 ? $" Weitere {reserved:F4} USD sind für laufende Läufe reserviert." : "";
|
||||
|
||||
private static string ReservedTokensHint(long reserved)
|
||||
=> reserved > 0 ? $" Weitere {reserved:N0} Tokens sind für laufende Läufe reserviert." : "";
|
||||
|
||||
/// <summary>
|
||||
/// Ohne Preisangaben ist die verbrauchte Summe unvollständig — sonst wirkte ein
|
||||
/// niedriger Wert wie ein noch offener Spielraum.
|
||||
|
||||
@@ -163,8 +163,11 @@ public sealed class AgentEngine : IAgentMessageRouter, Staging.IFrozenCallExecut
|
||||
return zu;
|
||||
|
||||
// Vor der ersten Anfrage prüfen — ein erschöpftes Budget soll gar nichts kosten.
|
||||
if (await CheckBudgetAsync(agentConfig, source ?? AuditSource.Direct, taskId, externalCt) is { } denied)
|
||||
// #716: Die Reservierung gilt, bis der Verbrauch erfasst ist (Ende der Methode).
|
||||
var (denied, reservation) = await CheckBudgetAsync(agentConfig, source ?? AuditSource.Direct, taskId, externalCt);
|
||||
if (denied is not null)
|
||||
return denied;
|
||||
using var budgetReservation = reservation;
|
||||
|
||||
using var lauf = TrackRun();
|
||||
|
||||
@@ -444,8 +447,10 @@ public sealed class AgentEngine : IAgentMessageRouter, Staging.IFrozenCallExecut
|
||||
if (RefuseWhileShuttingDown(agentConfig) is { } zu)
|
||||
return zu;
|
||||
|
||||
if (await CheckBudgetAsync(agentConfig, source, taskId, externalCt) is { } denied)
|
||||
var (denied, reservation) = await CheckBudgetAsync(agentConfig, source, taskId, externalCt);
|
||||
if (denied is not null)
|
||||
return denied;
|
||||
using var budgetReservation = reservation;
|
||||
|
||||
using var lauf = TrackRun();
|
||||
|
||||
@@ -1638,17 +1643,17 @@ public sealed class AgentEngine : IAgentMessageRouter, Staging.IFrozenCallExecut
|
||||
|
||||
/// <summary>
|
||||
/// Prüft das Tagesbudget. Gibt ein Ergebnis zurück, wenn der Lauf nicht stattfinden
|
||||
/// darf — sonst null.
|
||||
/// darf — sonst null und, bei gesetzter Grenze, die Reservierung des Laufs (#716).
|
||||
/// </summary>
|
||||
private async Task<AgentRunResult?> CheckBudgetAsync(
|
||||
private async Task<(AgentRunResult? Denied, IDisposable? Reservation)> CheckBudgetAsync(
|
||||
AgentConfig agentConfig, string? source, string? taskId, CancellationToken ct)
|
||||
{
|
||||
if (_budgetGuard is null)
|
||||
return null;
|
||||
return (null, null);
|
||||
|
||||
var status = await _budgetGuard.CheckAsync(agentConfig, InstanceBudget, ct);
|
||||
var (status, reservation) = await _budgetGuard.ReserveAsync(agentConfig, InstanceBudget, ct);
|
||||
if (status.IsAllowed)
|
||||
return null;
|
||||
return (null, reservation);
|
||||
|
||||
// #608: Ist ein Budget erschoepft, lehnt es jeden weiteren Lauf des Tages ab. Einmal je
|
||||
// Agent, Grenze und Tag als Warnung genuegt - danach steht jede Ablehnung als Beleg da.
|
||||
@@ -1676,7 +1681,7 @@ public sealed class AgentEngine : IAgentMessageRouter, Staging.IFrozenCallExecut
|
||||
await RecordReceiptAsync(Guid.NewGuid().ToString("N"), agentConfig, result, source, taskId);
|
||||
|
||||
OnRunCompleted?.Invoke(agentConfig.Model, result);
|
||||
return result;
|
||||
return (result, null);
|
||||
}
|
||||
|
||||
/// <summary>Welche Budgetgrenze je Agent zuletzt an welchem Tag als Warnung gemeldet wurde.</summary>
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
using ClawdDotNet.Core.Accounting;
|
||||
using ClawdDotNet.Core.Budget;
|
||||
using ClawdDotNet.Core.Config;
|
||||
using Shouldly;
|
||||
|
||||
namespace ClawdDotNet.Core.Tests.Budget;
|
||||
|
||||
/// <summary>
|
||||
/// #716 (the rest of #682): Runs in flight count against the budget with their estimate.
|
||||
///
|
||||
/// Before, every run was checked on its own against what was already recorded, so runs
|
||||
/// starting together all saw the same figure and all passed. #682 capped how many could
|
||||
/// start at once; the reservation closes the gap.
|
||||
/// </summary>
|
||||
public sealed class BudgetReservationTests
|
||||
{
|
||||
private static AgentConfig Agent(decimal dailyCost = 1.00m, long dailyTokens = 0)
|
||||
=> new() { AgentId = "worker", Budget = new AgentBudget { DailyCostUsd = dailyCost, DailyTokens = dailyTokens } };
|
||||
|
||||
/// <summary>Recent runs averaging 0.30 USD and 3,000 tokens.</summary>
|
||||
private static List<RunUsage> History(decimal cost = 0.30m, int tokens = 3_000)
|
||||
=> [.. Enumerable.Range(0, 5).Select(_ => new RunUsage
|
||||
{
|
||||
AgentId = "worker", PromptTokens = tokens, CostUsd = cost, CostIsKnown = true, Status = "Completed"
|
||||
})];
|
||||
|
||||
[Fact]
|
||||
public async Task Runs_in_flight_count_against_the_budget()
|
||||
{
|
||||
var guard = new BudgetGuard(new FakeUsage(usedCost: 0.60m, History()));
|
||||
|
||||
var (first, a) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);
|
||||
var (second, b) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);
|
||||
var (third, c) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);
|
||||
|
||||
first.IsAllowed.ShouldBeTrue("0.60 used, nothing in flight");
|
||||
second.IsAllowed.ShouldBeTrue("0.60 used + 0.30 in flight");
|
||||
third.IsAllowed.ShouldBeFalse("0.60 used + 0.60 in flight — before #716 this one passed too");
|
||||
third.Reason.ShouldNotBeNull().ShouldContain("reserviert");
|
||||
c.ShouldBeNull();
|
||||
|
||||
a!.Dispose();
|
||||
b!.Dispose();
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task A_finished_run_frees_its_reservation()
|
||||
{
|
||||
var guard = new BudgetGuard(new FakeUsage(usedCost: 0.60m, History()));
|
||||
|
||||
var (_, a) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);
|
||||
var (_, b) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);
|
||||
a!.Dispose();
|
||||
a.Dispose(); // twice is harmless
|
||||
|
||||
var (again, c) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);
|
||||
again.IsAllowed.ShouldBeTrue();
|
||||
|
||||
b!.Dispose();
|
||||
c!.Dispose();
|
||||
guard.Reserved("worker").ShouldBe(default(RunEstimate));
|
||||
guard.Reserved("").ShouldBe(default(RunEstimate));
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task A_single_run_is_not_refused_because_of_its_own_estimate()
|
||||
{
|
||||
// 0.95 of 1.00 used, an average run costs 0.30: a small task may still fit, and
|
||||
// refusing it on the average would be guessing.
|
||||
var guard = new BudgetGuard(new FakeUsage(usedCost: 0.95m, History()));
|
||||
|
||||
var (status, reservation) = await guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default);
|
||||
|
||||
status.IsAllowed.ShouldBeTrue();
|
||||
reservation!.Dispose();
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task Parallel_starts_cannot_all_pass_together()
|
||||
{
|
||||
var guard = new BudgetGuard(new FakeUsage(usedCost: 0m, History()));
|
||||
|
||||
var results = await Task.WhenAll(Enumerable.Range(0, 10)
|
||||
.Select(_ => guard.ReserveAsync(Agent(), InstanceBudget.Unlimited, default)));
|
||||
|
||||
// 0.00, 0.30, 0.60, 0.90 in flight pass; from 1.20 on the budget is spoken for.
|
||||
results.Count(r => r.Status.IsAllowed).ShouldBe(4);
|
||||
|
||||
foreach (var (_, reservation) in results)
|
||||
reservation?.Dispose();
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task The_token_limit_counts_reservations_too()
|
||||
{
|
||||
var guard = new BudgetGuard(new FakeUsage(usedCost: 0m, History(cost: 0m, tokens: 4_000)));
|
||||
var agent = Agent(dailyCost: 0, dailyTokens: 10_000);
|
||||
|
||||
var (a, ra) = await guard.ReserveAsync(agent, InstanceBudget.Unlimited, default);
|
||||
var (b, rb) = await guard.ReserveAsync(agent, InstanceBudget.Unlimited, default);
|
||||
var (c, _) = await guard.ReserveAsync(agent, InstanceBudget.Unlimited, default);
|
||||
var (d, _) = await guard.ReserveAsync(agent, InstanceBudget.Unlimited, default);
|
||||
|
||||
a.IsAllowed.ShouldBeTrue();
|
||||
b.IsAllowed.ShouldBeTrue();
|
||||
c.IsAllowed.ShouldBeTrue("8,000 of 10,000 reserved");
|
||||
d.IsAllowed.ShouldBeFalse("12,000 reserved");
|
||||
d.Limit.ShouldBe(BudgetLimit.AgentTokens);
|
||||
|
||||
ra!.Dispose();
|
||||
rb!.Dispose();
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task The_instance_limit_sees_the_runs_of_all_agents()
|
||||
{
|
||||
var guard = new BudgetGuard(new FakeUsage(usedCost: 0m, History()));
|
||||
var instance = new InstanceBudget { DailyCostUsd = 0.50m };
|
||||
var unlimitedAgent = Agent(dailyCost: 0);
|
||||
|
||||
var (a, _) = await guard.ReserveAsync(unlimitedAgent, instance, default);
|
||||
var (b, _) = await guard.ReserveAsync(unlimitedAgent, instance, default);
|
||||
var (c, _) = await guard.ReserveAsync(unlimitedAgent, instance, default);
|
||||
|
||||
a.IsAllowed.ShouldBeTrue();
|
||||
b.IsAllowed.ShouldBeTrue("0.30 in flight");
|
||||
c.IsAllowed.ShouldBeFalse("0.60 in flight over an instance limit of 0.50");
|
||||
c.Limit.ShouldBe(BudgetLimit.InstanceCost);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task Without_any_limit_nothing_is_reserved()
|
||||
{
|
||||
var guard = new BudgetGuard(new FakeUsage(usedCost: 0m, History()));
|
||||
|
||||
var (status, reservation) = await guard.ReserveAsync(Agent(dailyCost: 0), InstanceBudget.Unlimited, default);
|
||||
|
||||
status.IsAllowed.ShouldBeTrue();
|
||||
reservation.ShouldBeNull();
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task Runs_without_usage_do_not_drag_the_estimate_down()
|
||||
{
|
||||
var history = History();
|
||||
history.Add(new RunUsage { AgentId = "worker", Status = "BudgetExceeded" }); // 0 tokens
|
||||
var guard = new BudgetGuard(new FakeUsage(usedCost: 0m, history));
|
||||
|
||||
var estimate = await guard.EstimateRunAsync("worker", default);
|
||||
|
||||
estimate.CostUsd.ShouldBe(0.30m);
|
||||
estimate.Tokens.ShouldBe(3_000);
|
||||
}
|
||||
|
||||
private sealed class FakeUsage(decimal usedCost, IReadOnlyList<RunUsage> recent) : IUsageRepository
|
||||
{
|
||||
public Task RecordAsync(RunUsage usage, CancellationToken ct) => Task.CompletedTask;
|
||||
|
||||
public async Task<DailyUsage> GetDailyAsync(DateOnly date, string agentId, CancellationToken ct)
|
||||
{
|
||||
await Task.Yield(); // like the database: the caller does not get the answer synchronously
|
||||
return new DailyUsage(date, 0, 0, usedCost, true, 0);
|
||||
}
|
||||
|
||||
public Task<IReadOnlyList<RunUsage>> GetRecentAsync(int limit, CancellationToken ct)
|
||||
=> Task.FromResult(recent);
|
||||
|
||||
public Task<IReadOnlyList<RunUsage>> GetRecentForAgentAsync(string agentId, int limit, CancellationToken ct)
|
||||
=> Task.FromResult(recent);
|
||||
|
||||
public Task<int> PurgeBeforeAsync(DateOnly date, CancellationToken ct) => Task.FromResult(0);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user