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:
mika kuns
@@ -0,0 +1,165 @@
|
||||
using ClaudeDo.Data;
|
||||
using ClaudeDo.Data.Git;
|
||||
using ClaudeDo.Data.Models;
|
||||
using ClaudeDo.Data.Repositories;
|
||||
using ClaudeDo.Worker.Config;
|
||||
using ClaudeDo.Worker.Hub;
|
||||
using ClaudeDo.Worker.Runner;
|
||||
using ClaudeDo.Worker.Tests.Infrastructure;
|
||||
using ClaudeDo.Worker.Usage;
|
||||
using ClaudeDo.Worker.Usage.Interfaces;
|
||||
using Microsoft.Extensions.Logging.Abstractions;
|
||||
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;
|
||||
using Xunit;
|
||||
|
||||
namespace ClaudeDo.Worker.Tests.Runner;
|
||||
|
||||
/// Verifies TaskRunner persists raw token usage (input, output, cache-read, cache-write)
|
||||
/// aggregated from the session transcript via ITranscriptUsageReader, instead of the
|
||||
/// uncached-only counters the stream-json "result" event carries. See the bug report:
|
||||
/// task_runs.tokens_in was off by a factor of ~400,000 because it only read the API's
|
||||
/// per-call "input_tokens" field, ignoring cache_read/cache_creation.
|
||||
public sealed class RunUsagePersistenceTests : IDisposable
|
||||
{
|
||||
private readonly DbFixture _db = new();
|
||||
private readonly string _tempDir;
|
||||
private readonly WorkerConfig _cfg;
|
||||
private readonly FakeTranscriptUsageReader _reader = new();
|
||||
|
||||
public RunUsagePersistenceTests()
|
||||
{
|
||||
_tempDir = Path.Combine(Path.GetTempPath(), $"cd_runusage_{Guid.NewGuid():N}");
|
||||
Directory.CreateDirectory(_tempDir);
|
||||
_cfg = new WorkerConfig { SandboxRoot = _tempDir, LogRoot = _tempDir };
|
||||
}
|
||||
|
||||
public void Dispose() { _db.Dispose(); try { Directory.Delete(_tempDir, true); } catch { } }
|
||||
|
||||
private TaskRunner BuildRunner(IClaudeProcess claude, ITranscriptUsageReader? reader = null)
|
||||
{
|
||||
var dbFactory = _db.CreateFactory();
|
||||
var state = TaskStateServiceBuilder.Build(dbFactory).State;
|
||||
var wt = new WorktreeManager(new GitService(), dbFactory, _cfg, NullLogger<WorktreeManager>.Instance);
|
||||
return new TaskRunner(claude, dbFactory, new HubBroadcaster(new CapturingHubContext()), wt,
|
||||
new ClaudeArgsBuilder(), _cfg, NullLogger<TaskRunner>.Instance, state, new TaskRunTokenRegistry(),
|
||||
new AttachmentStore(), new FakeSessionSkillSeeder(), reader ?? _reader);
|
||||
}
|
||||
|
||||
private async Task SeedTaskAsync(string taskId, TaskStatus status = TaskStatus.Idle)
|
||||
{
|
||||
using var ctx = _db.CreateContext();
|
||||
var listId = Guid.NewGuid().ToString();
|
||||
ctx.Lists.Add(new ListEntity { Id = listId, Name = "L", WorkingDir = null, CreatedAt = DateTime.UtcNow });
|
||||
ctx.Tasks.Add(new TaskEntity
|
||||
{
|
||||
Id = taskId, ListId = listId, Title = "Task", Status = status, CreatedAt = DateTime.UtcNow,
|
||||
});
|
||||
await ctx.SaveChangesAsync();
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task First_run_persists_all_four_token_fields_from_the_transcript_reader()
|
||||
{
|
||||
const string taskId = "t1";
|
||||
await SeedTaskAsync(taskId);
|
||||
_reader.SetTotals("sess-a", new SessionUsageTotals(10, 20, 300, 5));
|
||||
|
||||
var fake = new FakeClaudeProcess((_, _, _, _, _) =>
|
||||
Task.FromResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok", SessionId = "sess-a" }));
|
||||
var runner = BuildRunner(fake);
|
||||
|
||||
using (var ctx = _db.CreateContext())
|
||||
await runner.RunAsync((await new TaskRepository(ctx).GetByIdAsync(taskId))!, "slot-1", CancellationToken.None);
|
||||
|
||||
using var verify = _db.CreateContext();
|
||||
var run = await new TaskRunRepository(verify).GetLatestByTaskIdAsync(taskId);
|
||||
Assert.Equal(10, run!.TokensIn);
|
||||
Assert.Equal(20, run.TokensOut);
|
||||
Assert.Equal(300, run.CacheReadTokens);
|
||||
Assert.Equal(5, run.CacheWriteTokens);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task Resumed_run_on_the_same_session_persists_only_the_delta()
|
||||
{
|
||||
const string taskId = "t2";
|
||||
await SeedTaskAsync(taskId, TaskStatus.WaitingForReview);
|
||||
_reader.SetTotals("sess-b", new SessionUsageTotals(10, 20, 100, 0));
|
||||
|
||||
var fake = new FakeClaudeProcess((_, _, _, _, _) =>
|
||||
Task.FromResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok", SessionId = "sess-b" }));
|
||||
var runner = BuildRunner(fake);
|
||||
|
||||
using (var ctx = _db.CreateContext())
|
||||
await runner.RunAsync((await new TaskRepository(ctx).GetByIdAsync(taskId))!, "slot-1", CancellationToken.None);
|
||||
|
||||
// Cumulative transcript totals grow after the second (resumed) turn.
|
||||
_reader.SetTotals("sess-b", new SessionUsageTotals(30, 50, 250, 10));
|
||||
await runner.ContinueAsync(taskId, "follow up", "slot-1", CancellationToken.None);
|
||||
|
||||
using var verify = _db.CreateContext();
|
||||
var runs = await new TaskRunRepository(verify).GetByTaskIdAsync(taskId);
|
||||
Assert.Equal(2, runs.Count);
|
||||
var run2 = runs.Single(r => r.RunNumber == 2);
|
||||
Assert.Equal(20, run2.TokensIn);
|
||||
Assert.Equal(30, run2.TokensOut);
|
||||
Assert.Equal(150, run2.CacheReadTokens);
|
||||
Assert.Equal(10, run2.CacheWriteTokens);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task Missing_transcript_totals_leave_token_fields_null_but_run_still_succeeds()
|
||||
{
|
||||
const string taskId = "t3";
|
||||
await SeedTaskAsync(taskId);
|
||||
// No totals registered for "sess-missing" -> reader returns null.
|
||||
|
||||
var fake = new FakeClaudeProcess((_, _, _, _, _) =>
|
||||
Task.FromResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok", SessionId = "sess-missing" }));
|
||||
var runner = BuildRunner(fake);
|
||||
|
||||
using (var ctx = _db.CreateContext())
|
||||
await runner.RunAsync((await new TaskRepository(ctx).GetByIdAsync(taskId))!, "slot-1", CancellationToken.None);
|
||||
|
||||
using var verify = _db.CreateContext();
|
||||
var run = await new TaskRunRepository(verify).GetLatestByTaskIdAsync(taskId);
|
||||
Assert.Null(run!.TokensIn);
|
||||
Assert.Null(run.TokensOut);
|
||||
Assert.Null(run.CacheReadTokens);
|
||||
Assert.Null(run.CacheWriteTokens);
|
||||
Assert.Equal(0, run.ExitCode);
|
||||
|
||||
var task = await new TaskRepository(verify).GetByIdAsync(taskId);
|
||||
Assert.Equal(TaskStatus.WaitingForReview, task!.Status);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task A_failing_usage_reader_does_not_fail_the_run()
|
||||
{
|
||||
const string taskId = "t4";
|
||||
await SeedTaskAsync(taskId);
|
||||
|
||||
var fake = new FakeClaudeProcess((_, _, _, _, _) =>
|
||||
Task.FromResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok", SessionId = "sess-throws" }));
|
||||
var runner = BuildRunner(fake, new ThrowingUsageReader());
|
||||
|
||||
using (var ctx = _db.CreateContext())
|
||||
await runner.RunAsync((await new TaskRepository(ctx).GetByIdAsync(taskId))!, "slot-1", CancellationToken.None);
|
||||
|
||||
using var verify = _db.CreateContext();
|
||||
var run = await new TaskRunRepository(verify).GetLatestByTaskIdAsync(taskId);
|
||||
Assert.Null(run!.TokensIn);
|
||||
var task = await new TaskRepository(verify).GetByIdAsync(taskId);
|
||||
Assert.Equal(TaskStatus.WaitingForReview, task!.Status);
|
||||
}
|
||||
|
||||
private sealed class ThrowingUsageReader : ITranscriptUsageReader
|
||||
{
|
||||
public Task<IReadOnlyList<UsageAggregateRow>> ReadAsync(
|
||||
DateOnly start, DateOnly end, CancellationToken ct = default) =>
|
||||
throw new IOException("boom");
|
||||
|
||||
public Task<SessionUsageTotals?> ReadSessionTotalsAsync(string sessionId, CancellationToken ct = default) =>
|
||||
throw new IOException("boom");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user