ClawdDotNet/tests/ClawdDotNet.Core.Tests/Budget/BudgetHeadroomTests.cs

using ClawdDotNet.Core.Accounting;
using ClawdDotNet.Core.Budget;
using ClawdDotNet.Core.Config;
using ClawdDotNet.Core.Tasks;
using Shouldly;

namespace ClawdDotNet.Core.Tests.Budget;

/// <summary>
/// #716, stage 3: late work — carried-over tasks and catch-ups — only uses budget the rest
/// of the day does not need. The new day has right of way.
/// </summary>
public sealed class BudgetHeadroomTests
{
    // 2027-03-10 14:00 UTC; the test clock's local zone is UTC, so the day ends at 24:00 UTC.
    private static readonly DateTime Now = new(2027, 3, 10, 14, 0, 0, DateTimeKind.Utc);
    private static readonly DateTime CreatedEarlier = new(2027, 3, 1, 0, 0, 0, DateTimeKind.Utc);

    private static AgentConfig Agent(decimal dailyCost = 1.00m)
        => new() { AgentId = "worker", Budget = new AgentBudget { DailyCostUsd = dailyCost } };

    /// <summary>Runs averaging 0.05 USD.</summary>
    private static List<RunUsage> History(decimal cost = 0.05m)
        => [.. Enumerable.Range(0, 5).Select(_ => new RunUsage { AgentId = "worker", PromptTokens = 1_000, CostUsd = cost, CostIsKnown = true })];

    private static TaskItem Cron(string id, string expression, string? lastOccurrence = null) => new()
    {
        Id = id, Title = id, Assignee = "@worker", CreatedAt = CreatedEarlier, LastOccurrence = lastOccurrence,
        When = new TaskWhen { Kind = TaskWhenKind.Cron, Value = expression, TimeZone = "UTC" }
    };

    private static TaskItem Carried(string id) => new()
    {
        Id = id, Title = id, Assignee = "@worker", CreatedAt = CreatedEarlier, BudgetMissCount = 1
    };

    private static BudgetHeadroom Headroom(decimal used, AgentConfig agent, InstanceBudget? instance = null, decimal perRun = 0.05m)
        => new(new BudgetGuard(new FakeUsage(used, History(perRun))),
            () => [agent], () => instance ?? InstanceBudget.Unlimited, new UtcClock(Now));

    // ═══════════════════ counting the rest of the day ═══════════════════

    [Fact]
    public void Hourly_from_14_30_until_midnight_is_nine_more_runs()
    {
        // 15:00 … 23:00. The 14:00 run is done (marker), so nothing is due right now.
        var task = Cron("t", "0 * * * *", lastOccurrence: "2027-03-10T14:00:00.0000000Z");

        TaskSchedule.OccurrencesBetween(task, Now.AddMinutes(30), Now.Date.AddDays(1)).ShouldBe(9);
    }

    [Fact]
    public void A_run_due_right_now_counts_as_well()
    {
        // Last run 13:00, so the 14:00 run is due now; then 15:00 … 23:00.
        var task = Cron("t", "0 * * * *", lastOccurrence: "2027-03-10T13:00:00.0000000Z");

        TaskSchedule.OccurrencesBetween(task, Now.AddMinutes(1), Now.Date.AddDays(1)).ShouldBe(10);
    }

    [Fact]
    public void A_one_time_task_counts_once_until_it_ran()
    {
        var open = new TaskItem { Id = "t", CreatedAt = CreatedEarlier };
        var done = open with { LastOccurrence = "2027-03-10T09:00:00.0000000Z" };

        TaskSchedule.OccurrencesBetween(open, Now, Now.AddHours(1)).ShouldBe(1);
        TaskSchedule.OccurrencesBetween(done, Now, Now.AddHours(1)).ShouldBe(0);
    }

    [Fact]
    public void An_interval_counts_its_remaining_ticks()
    {
        var task = new TaskItem
        {
            Id = "t", CreatedAt = CreatedEarlier, LastOccurrence = "2027-03-10T13:45:00.0000000Z",
            When = new TaskWhen { Kind = TaskWhenKind.Every, Value = "30m", TimeZone = "UTC" }
        };

        // 14:15, 14:45, 15:15, 15:45 before 16:00.
        TaskSchedule.OccurrencesBetween(task, Now, Now.AddHours(2)).ShouldBe(4);
    }

    // ═══════════════════ headroom ═══════════════════

    [Fact]
    public async Task Late_work_fits_when_the_rest_of_the_day_leaves_room()
    {
        // Limit 1.00, used 0.30, margin 0.10. Rest of the day: 10 hourly runs à 0.05 = 0.50.
        // Left: 1.00 - 0.30 - 0.50 - 0.10 = 0.10 ≥ 0.05 for the late run.
        var hourly = Cron("t-hourly", "0 * * * *", lastOccurrence: "2027-03-10T13:00:00.0000000Z");
        var late = Carried("t-late");

        var room = await Headroom(used: 0.30m, Agent()).HasRoomAsync(late, [hourly, late], Now, default);

        room.ShouldBeTrue();
    }

    [Fact]
    public async Task Late_work_waits_when_the_rest_of_the_day_needs_the_budget()
    {
        // Same, but 0.40 used: 1.00 - 0.40 - 0.50 - 0.10 = 0.00 < 0.05.
        var hourly = Cron("t-hourly", "0 * * * *", lastOccurrence: "2027-03-10T13:00:00.0000000Z");
        var late = Carried("t-late");

        var room = await Headroom(used: 0.40m, Agent()).HasRoomAsync(late, [hourly, late], Now, default);

        room.ShouldBeFalse("the day's own schedule has right of way");
    }

    [Fact]
    public async Task Other_late_work_does_not_count_as_the_days_demand()
    {
        // Ten carried-over tasks do not reserve budget against each other — they are what
        // waits for headroom, not what headroom is kept for.
        var late = Carried("t-late");
        var others = Enumerable.Range(0, 10).Select(i => Carried($"t-other-{i}")).ToList();

        var room = await Headroom(used: 0.50m, Agent()).HasRoomAsync(late, [late, .. others], Now, default);

        room.ShouldBeTrue();
    }

    [Fact]
    public async Task The_instance_limit_counts_the_schedules_of_all_agents()
    {
        var agent = Agent(dailyCost: 0); // no limit of its own
        var instance = new InstanceBudget { DailyCostUsd = 1.00m };
        var hourly = Cron("t-hourly", "0 * * * *", lastOccurrence: "2027-03-10T13:00:00.0000000Z");
        var late = Carried("t-late");

        (await Headroom(used: 0.30m, agent, instance).HasRoomAsync(late, [hourly, late], Now, default)).ShouldBeTrue();
        (await Headroom(used: 0.40m, agent, instance).HasRoomAsync(late, [hourly, late], Now, default)).ShouldBeFalse();
    }

    [Fact]
    public async Task Without_any_limit_there_is_always_room()
    {
        var late = Carried("t-late");

        (await Headroom(used: 99m, Agent(dailyCost: 0)).HasRoomAsync(late, [late], Now, default)).ShouldBeTrue();
    }

    // ═══════════════════ the day ahead ═══════════════════

    [Fact]
    public async Task The_outlook_names_the_hour_the_budget_runs_out()
    {
        // Limit 1.00, used 0.60, runs at 14:00 (due now), 15:00 … 23:00 à 0.05. Before the
        // 22:00 run 0.60 + 8 × 0.05 = 1.00 is used — from there on everything is refused.
        var hourly = Cron("t-hourly", "0 * * * *", lastOccurrence: "2027-03-10T13:00:00.0000000Z");

        var shortfalls = await Headroom(used: 0.60m, Agent()).ShortfallsAsync([hourly], Now, default);

        var shortfall = shortfalls.ShouldHaveSingleItem();
        shortfall.AgentId.ShouldBe("worker");
        shortfall.ExhaustsAtUtc.ShouldBe(new DateTime(2027, 3, 10, 22, 0, 0, DateTimeKind.Utc));
        shortfall.RunsRefused.ShouldBe(2);
        shortfall.RunsPlanned.ShouldBe(10);
    }

    [Fact]
    public async Task A_day_that_fits_has_no_shortfall()
    {
        var hourly = Cron("t-hourly", "0 * * * *", lastOccurrence: "2027-03-10T13:00:00.0000000Z");

        (await Headroom(used: 0.30m, Agent()).ShortfallsAsync([hourly], Now, default)).ShouldBeEmpty();
    }

    [Fact]
    public async Task Late_work_does_not_make_the_day_fall_short()
    {
        // Fifty carried-over tasks only ever run in headroom — they cannot be the reason the
        // day's own runs are refused.
        var hourly = Cron("t-hourly", "0 * * * *", lastOccurrence: "2027-03-10T13:00:00.0000000Z");
        var late = Enumerable.Range(0, 50).Select(i => Carried($"t-late-{i}")).ToList();

        (await Headroom(used: 0.30m, Agent()).ShortfallsAsync([hourly, .. late], Now, default)).ShouldBeEmpty();
    }

    [Fact]
    public async Task The_instance_limit_gets_its_own_forecast()
    {
        var hourly = Cron("t-hourly", "0 * * * *", lastOccurrence: "2027-03-10T13:00:00.0000000Z");
        var instance = new InstanceBudget { DailyCostUsd = 1.00m };

        var shortfalls = await Headroom(used: 0.60m, Agent(dailyCost: 0), instance).ShortfallsAsync([hourly], Now, default);

        shortfalls.ShouldHaveSingleItem().AgentId.ShouldBe("", "empty means: the instance");
    }

    // ═══════════════════ test doubles ═══════════════════

    private sealed class UtcClock(DateTime now) : TimeProvider
    {
        public override DateTimeOffset GetUtcNow() => new(now, TimeSpan.Zero);
        public override TimeZoneInfo LocalTimeZone => TimeZoneInfo.Utc;
    }

    private sealed class FakeUsage(decimal usedCost, IReadOnlyList<RunUsage> recent) : IUsageRepository
    {
        public Task RecordAsync(RunUsage usage, CancellationToken ct) => Task.CompletedTask;

        public Task<DailyUsage> GetDailyAsync(DateOnly date, string agentId, CancellationToken ct)
            => Task.FromResult(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);
    }
}