ClawdDotNet/src/ClawdDotNet.Core/Budget/TaskCostEstimator.cs

using ClawdDotNet.Core.Audit;
using ClawdDotNet.Core.Tasks;

namespace ClawdDotNet.Core.Budget;

/// <summary>#717: Where an estimate came from — from the most to the least reliable.</summary>
public enum EstimateSource
{
    /// <summary>The task's own runs, measured.</summary>
    Measured,

    /// <summary>Runs of its siblings — the other parts of the same split task.</summary>
    Siblings,

    /// <summary>What the task's author wrote down. A guess, and shown as one.</summary>
    Stated,

    /// <summary>The average run of the agent — says nothing about this task in particular.</summary>
    AgentAverage,

    /// <summary>Nothing to go on.</summary>
    None
}

/// <summary>#717: What one run of a task is expected to cost, and how much that is worth.</summary>
public sealed record TaskEstimate(RunEstimate Run, EstimateSource Source, int Samples)
{
    /// <summary>Stated or averaged over everything the agent does — not measured on this work.</summary>
    public bool IsGuess => Source is EstimateSource.Stated or EstimateSource.AgentAverage or EstimateSource.None;

    public static TaskEstimate Nothing { get; } = new(default, EstimateSource.None, 0);
}

/// <summary>
/// #717: Estimates a run of one particular task — honestly.
///
/// <para>Costs cannot be known in advance, and models judge their own effort badly. So the
/// code estimates from history where there is any, and says when it has none:</para>
/// <list type="number">
/// <item>the task's own runs (a fill task that already ran part of the way, or ran before),</item>
/// <item>its siblings — the parts of the same split task, measured as soon as the first one ran,</item>
/// <item>the author's estimate (<see cref="TaskItem.EstimateUsd"/>), marked as a guess,</item>
/// <item>the agent's average run, also a guess.</item>
/// </list>
/// <para>After the first run the measured value replaces the guess — nothing has to be
/// updated for that, the next estimate simply finds the run.</para>
/// </summary>
public sealed class TaskCostEstimator
{
    /// <summary>How many runs are averaged — recent enough to follow a change, enough to smooth one outlier.</summary>
    public const int Sample = 10;

    private readonly IAuditRepository? _audit;
    private readonly ITaskRepository _tasks;
    private readonly BudgetGuard _guard;

    public TaskCostEstimator(IAuditRepository? audit, ITaskRepository tasks, BudgetGuard guard)
    {
        _audit = audit;
        _tasks = tasks;
        _guard = guard;
    }

    public async Task<TaskEstimate> EstimateAsync(TaskItem task, CancellationToken ct)
    {
        if (_audit is not null)
        {
            var own = Measured(await _audit.ListReceiptsForTaskAsync(task.Id, ct));
            if (own.Count > 0)
                return Average(own, EstimateSource.Measured);

            if (!string.IsNullOrWhiteSpace(task.ParentId))
            {
                var siblingRuns = new List<RunReceipt>();
                foreach (var sibling in await _tasks.ListChildrenAsync(task.ParentId, ct))
                {
                    if (sibling.Id == task.Id)
                        continue;
                    siblingRuns.AddRange(Measured(await _audit.ListReceiptsForTaskAsync(sibling.Id, ct)));
                }
                if (siblingRuns.Count > 0)
                    return Average(siblingRuns.OrderByDescending(r => r.OccurredAt).Take(Sample).ToList(), EstimateSource.Siblings);
            }
        }

        var agentId = TaskAssignee.AgentId(task.Assignee);
        var agentAverage = agentId.Length == 0 ? default : await _guard.EstimateRunAsync(agentId, ct);

        if (task.EstimateUsd is { } stated)
        {
            // Tokens follow the agent's own price per token, where it has one — a token limit
            // needs a token figure, and the author only wrote down money.
            var tokens = agentAverage.CostUsd > 0
                ? (long)Math.Ceiling(stated / agentAverage.CostUsd * agentAverage.Tokens)
                : 0;
            return new TaskEstimate(new RunEstimate(stated, tokens), EstimateSource.Stated, 0);
        }

        return agentAverage == default
            ? TaskEstimate.Nothing
            : new TaskEstimate(agentAverage, EstimateSource.AgentAverage, 0);
    }

    /// <summary>Runs that used something — a refusal or an abort before the first call says nothing about cost.</summary>
    private static List<RunReceipt> Measured(IReadOnlyList<RunReceipt> receipts)
        => receipts.Where(r => r.PromptTokens + r.CompletionTokens > 0)
            .OrderByDescending(r => r.OccurredAt)
            .Take(Sample)
            .ToList();

    private static TaskEstimate Average(IReadOnlyList<RunReceipt> runs, EstimateSource source)
        => new(new RunEstimate(
                runs.Average(r => r.CostUsd),
                (long)Math.Ceiling(runs.Average(r => (double)(r.PromptTokens + r.CompletionTokens)))),
            source, runs.Count);
}