diff --git a/src/ClawdDotNet.Core/Accounting/SqliteUsageRepository.cs b/src/ClawdDotNet.Core/Accounting/SqliteUsageRepository.cs index a63364d..40df9f4 100644 --- a/src/ClawdDotNet.Core/Accounting/SqliteUsageRepository.cs +++ b/src/ClawdDotNet.Core/Accounting/SqliteUsageRepository.cs @@ -125,6 +125,29 @@ public sealed class SqliteUsageRepository : IUsageRepository return results; } + public async Task> 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(); + await using var reader = await cmd.ExecuteReaderAsync(ct); + while (await reader.ReadAsync(ct)) + results.Add(Read(reader)); + + return results; + } + public Task PurgeBeforeAsync(DateOnly date, CancellationToken ct) => _storage.WriteAsync(async conn => { diff --git a/src/ClawdDotNet.Core/Accounting/UsageModels.cs b/src/ClawdDotNet.Core/Accounting/UsageModels.cs index b60a975..f66d403 100644 --- a/src/ClawdDotNet.Core/Accounting/UsageModels.cs +++ b/src/ClawdDotNet.Core/Accounting/UsageModels.cs @@ -53,6 +53,14 @@ public interface IUsageRepository Task> GetRecentAsync(int limit, CancellationToken ct); + /// + /// #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. + /// + Task> GetRecentForAgentAsync(string agentId, int limit, CancellationToken ct) + => Task.FromResult>([]); + /// Entfernt Einträge älter als das angegebene Datum. Gibt die Anzahl zurück. Task PurgeBeforeAsync(DateOnly date, CancellationToken ct); } diff --git a/src/ClawdDotNet.Core/Budget/BudgetGuard.cs b/src/ClawdDotNet.Core/Budget/BudgetGuard.cs index 2ecac6b..6419b60 100644 --- a/src/ClawdDotNet.Core/Budget/BudgetGuard.cs +++ b/src/ClawdDotNet.Core/Budget/BudgetGuard.cs @@ -13,6 +13,13 @@ public enum BudgetLimit InstanceTokens } +/// #716: What a run is expected to use — the agent's average over recent runs. +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 _now; + /// How many recent runs of an agent make up its per-run estimate (#716). + public const int EstimateSample = 20; + + /// + /// #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. + /// + private readonly SemaphoreSlim _reserveGate = new(1, 1); + + private readonly Lock _reservedLock = new(); + private readonly Dictionary _reservedByAgent = new(StringComparer.OrdinalIgnoreCase); + private RunEstimate _reservedInstance; + public BudgetGuard(IUsageRepository usage, Func? now = null) { _usage = usage; @@ -43,8 +63,112 @@ public sealed class BudgetGuard /// /// Prüft, ob der Agent heute noch laufen darf. /// - public async Task CheckAsync( + public Task CheckAsync( AgentConfig agent, InstanceBudget instanceBudget, CancellationToken ct) + => CheckCoreAsync(agent, instanceBudget, default, default, ct); + + /// + /// #716: Checks like 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. + /// + /// What changes. 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. + /// + /// What does not change. 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. + /// + 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(); + } + } + + /// What is currently reserved by runs in flight — for the agent, or with an + /// empty id for the whole instance. + public RunEstimate Reserved(string agentId) + { + lock (_reservedLock) + return agentId.Length == 0 ? _reservedInstance : _reservedByAgent.GetValueOrDefault(agentId); + } + + /// + /// 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. + /// + public async Task 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 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; + /// #716: Without this, a refusal at "3.20 of 5.00 USD used" would look wrong. + 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." : ""; + /// /// Ohne Preisangaben ist die verbrauchte Summe unvollständig — sonst wirkte ein /// niedriger Wert wie ein noch offener Spielraum. diff --git a/src/ClawdDotNet.Core/Engine/AgentEngine.cs b/src/ClawdDotNet.Core/Engine/AgentEngine.cs index b38a0c3..cc8389b 100644 --- a/src/ClawdDotNet.Core/Engine/AgentEngine.cs +++ b/src/ClawdDotNet.Core/Engine/AgentEngine.cs @@ -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 /// /// 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). /// - private async Task 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); } /// Welche Budgetgrenze je Agent zuletzt an welchem Tag als Warnung gemeldet wurde. diff --git a/tests/ClawdDotNet.Core.Tests/Budget/BudgetReservationTests.cs b/tests/ClawdDotNet.Core.Tests/Budget/BudgetReservationTests.cs new file mode 100644 index 0000000..69418ea --- /dev/null +++ b/tests/ClawdDotNet.Core.Tests/Budget/BudgetReservationTests.cs @@ -0,0 +1,173 @@ +using ClawdDotNet.Core.Accounting; +using ClawdDotNet.Core.Budget; +using ClawdDotNet.Core.Config; +using Shouldly; + +namespace ClawdDotNet.Core.Tests.Budget; + +/// +/// #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. +/// +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 } }; + + /// Recent runs averaging 0.30 USD and 3,000 tokens. + private static List 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 recent) : IUsageRepository + { + public Task RecordAsync(RunUsage usage, CancellationToken ct) => Task.CompletedTask; + + public async Task 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> GetRecentAsync(int limit, CancellationToken ct) + => Task.FromResult(recent); + + public Task> GetRecentForAgentAsync(string agentId, int limit, CancellationToken ct) + => Task.FromResult(recent); + + public Task PurgeBeforeAsync(DateOnly date, CancellationToken ct) => Task.FromResult(0); + } +}