fix(worker): record real raw token usage per run, not the uncached remainder
task_runs.tokens_in only ever held the API's uncached "input_tokens" field (off by a factor of ~400,000 on a resumed session), and tokens_out summed only the last result event instead of the whole session. TaskRunner now reads each run's cache-read/cache-write/input/output totals from the session transcript via a new ITranscriptUsageReader.ReadSessionTotalsAsync, storing the delta against prior runs on the same session so a --resume doesn't double-count. New task_runs.cache_read_tokens/cache_write_tokens columns; the Session tab now shows the raw total (what actually counts against the 5h/7d usage limit) with a breakdown tooltip.
This commit is contained in:
@@ -6,6 +6,8 @@ using ClaudeDo.Worker.Config;
|
||||
using ClaudeDo.Worker.Hub;
|
||||
using ClaudeDo.Worker.Skills;
|
||||
using ClaudeDo.Worker.State;
|
||||
using ClaudeDo.Worker.Usage;
|
||||
using ClaudeDo.Worker.Usage.Interfaces;
|
||||
using Microsoft.EntityFrameworkCore;
|
||||
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;
|
||||
|
||||
@@ -24,6 +26,7 @@ public sealed class TaskRunner
|
||||
private readonly TaskRunTokenRegistry _tokens;
|
||||
private readonly AttachmentStore _attachments;
|
||||
private readonly ISessionSkillSeeder _skillSeeder;
|
||||
private readonly ITranscriptUsageReader _usageReader;
|
||||
|
||||
public TaskRunner(
|
||||
IClaudeProcess claude,
|
||||
@@ -36,7 +39,8 @@ public sealed class TaskRunner
|
||||
ITaskStateService state,
|
||||
TaskRunTokenRegistry tokens,
|
||||
AttachmentStore attachments,
|
||||
ISessionSkillSeeder skillSeeder)
|
||||
ISessionSkillSeeder skillSeeder,
|
||||
ITranscriptUsageReader usageReader)
|
||||
{
|
||||
_claude = claude;
|
||||
_dbFactory = dbFactory;
|
||||
@@ -49,6 +53,7 @@ public sealed class TaskRunner
|
||||
_tokens = tokens;
|
||||
_attachments = attachments;
|
||||
_skillSeeder = skillSeeder;
|
||||
_usageReader = usageReader;
|
||||
}
|
||||
|
||||
public async Task RunAsync(TaskEntity task, string slot, CancellationToken ct, bool alreadyClaimed = false)
|
||||
@@ -361,8 +366,8 @@ public sealed class TaskRunner
|
||||
run.ErrorMarkdown = result.ErrorMarkdown;
|
||||
run.ExitCode = result.ExitCode;
|
||||
run.TurnCount = result.TurnCount;
|
||||
run.TokensIn = result.TokensIn;
|
||||
run.TokensOut = result.TokensOut;
|
||||
if (result.SessionId is not null)
|
||||
await ApplyUsageAsync(run, taskId, result.SessionId);
|
||||
run.FinishedAt = DateTime.UtcNow;
|
||||
|
||||
using (var context = _dbFactory.CreateDbContext())
|
||||
@@ -394,6 +399,40 @@ public sealed class TaskRunner
|
||||
}
|
||||
}
|
||||
|
||||
/// Populates the run's raw token fields from the session transcript (input, output,
|
||||
/// cache-read, cache-write — the API's "input_tokens" alone is only the uncached
|
||||
/// remainder and undercounts the real prompt size by orders of magnitude). A resumed
|
||||
/// session's transcript is cumulative, so the delta against prior runs sharing the same
|
||||
/// SessionId is stored, not the running total. Any failure here (missing/unreadable
|
||||
/// transcript) leaves the fields null and must never fail the run itself.
|
||||
private async Task ApplyUsageAsync(TaskRunEntity run, string taskId, string sessionId)
|
||||
{
|
||||
try
|
||||
{
|
||||
var totals = await _usageReader.ReadSessionTotalsAsync(sessionId, CancellationToken.None);
|
||||
if (totals is null) return;
|
||||
|
||||
List<TaskRunEntity> priorRuns;
|
||||
using (var context = _dbFactory.CreateDbContext())
|
||||
priorRuns = await new TaskRunRepository(context).GetByTaskIdAsync(taskId, CancellationToken.None);
|
||||
|
||||
var sameSession = priorRuns.Where(r => r.Id != run.Id && r.SessionId == sessionId).ToList();
|
||||
var priorInput = sameSession.Sum(r => (long)(r.TokensIn ?? 0));
|
||||
var priorOutput = sameSession.Sum(r => (long)(r.TokensOut ?? 0));
|
||||
var priorCacheRead = sameSession.Sum(r => (long)(r.CacheReadTokens ?? 0));
|
||||
var priorCacheWrite = sameSession.Sum(r => (long)(r.CacheWriteTokens ?? 0));
|
||||
|
||||
run.TokensIn = (int)Math.Max(0, totals.InputTokens - priorInput);
|
||||
run.TokensOut = (int)Math.Max(0, totals.OutputTokens - priorOutput);
|
||||
run.CacheReadTokens = (int)Math.Max(0, totals.CacheReadTokens - priorCacheRead);
|
||||
run.CacheWriteTokens = (int)Math.Max(0, totals.CacheCreationTokens - priorCacheWrite);
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
_logger.LogWarning(ex, "Failed to read session usage totals for task {TaskId}, session {SessionId}", taskId, sessionId);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task HandleSuccess(TaskEntity task, ListEntity list, string slot, WorktreeContext? wtCtx, RunResult result, CancellationToken ct)
|
||||
{
|
||||
if (wtCtx is not null)
|
||||
|
||||
Reference in New Issue
Block a user