ClaudeDo/tests/ClaudeDo.Worker.Tests/Runner/RunUsagePersistenceTests.cs

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");
    }
}