fix(review): sum live tokens across a collapsed agent's workers · Entire

fix(review): sum live tokens across a collapsed agent's workers

ce3d125·

peyton-alt·3d ago·2 files·+82 added/-1 removed

Follow-on to the one-row-per-agent collapse: #1666 (merged) streams cumulative token counts per worker, and the TUI row overwrites on each Tokens event (CU2). With two skill-workers folded into one agent row, their independent cumulative counts overwrote each other — the live number bounced between skills instead of showing the agent total. (The final summary was already summed; only the live display was wrong.)

The sink now tracks each worker's latest cumulative Tokens and forwards the per-agent SUM on every Tokens event, so the row's overwrite shows a correct running total. Single-worker rows pass through unchanged. Runs on the serial dispatch goroutine, so no locking; race-clean.

Co-Authored-By: Claude Fable 5 noreply@anthropic.com

Sessions

01KXH2J8TTE8GJYBQE47W4BHKXView transcript

[?
test(review): pin that codex $name skills survive the legacy repair + native invocationClaude Code·4 steps](/content/gh/entireio/cli/session/99f0d90a-6e31-4c60-ba62-aa2dccd29033#timeline-01KXH2J8TTE8GJYBQE47W4BHKX/index.html)

Changes

2

71 unmodified lines

72
73
74
75
76
77
78
79
80
81
82
83
3 unmodified lines

87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
212 unmodified lines

324
325
326
327
328
329
330
331
332
333
334
335
303
336
337
338
339

71 unmodified lines

// and the summary pass through unchanged.
    rowOrder      []string
    workerToAgent map[string]string
    rowWorkers    map[string][]string // agent row → its worker labels
    // workerTokens holds each worker's latest cumulative Tokens so a
    // collapsed row can display the per-agent SUM live (workers stream
    // cumulative counts independently; overwriting would bounce between
    // them). Written and read only from the serial dispatch goroutine.
    workerTokens map[string]reviewtypes.Tokens
    }
// groupWorkers configures one-row-per-agent collapsing: rowOrder is the
3 unmodified lines

func (s *TUISink) groupWorkers(rowOrder []string, workerToAgent map[string]string) {
    s.rowOrder = rowOrder
    s.workerToAgent = workerToAgent
    s.rowWorkers = make(map[string][]string, len(rowOrder))
    for worker, row := range workerToAgent {
        s.rowWorkers[row] = append(s.rowWorkers[row], worker)
    }
    s.workerTokens = make(map[string]reviewtypes.Tokens, len(workerToAgent))
}

// agentRowTokens records worker's latest cumulative Tokens and returns the
// summed Tokens across every worker folded into agentRow. Called only from
// the serial dispatch goroutine (CU4 contract), so it needs no lock.
func (s *TUISink) agentRowTokens(worker, agentRow string, tk reviewtypes.Tokens) reviewtypes.Tokens {
    s.workerTokens[worker] = tk
    var sum reviewtypes.Tokens
    for _, w := range s.rowWorkers[agentRow] {
        wt := s.workerTokens[w]
        sum.In += wt.In
        sum.Out += wt.Out
    }
    return sum
}

// agentRowFor resolves a worker label to its agent row, passing through any
212 unmodified lines

if !ok {
        return
    }
    row := s.agentRowFor(agent)
    // A collapsed row shows the SUM of its workers' cumulative token counts;
    // forwarding a single worker's cumulative value would bounce the live
    // number between siblings (row.tokens overwrites, per CU2). Ungrouped
    // (single-worker) rows pass through unchanged.
    if tk, ok := ev.(reviewtypes.Tokens); ok && s.workerTokens != nil {
        ev = s.agentRowTokens(agent, row, tk)
    }
    select {
    case s.msgs <- agentEventMsg{agent: s.agentRowFor(agent), ev: ev}:
    case s.msgs <- agentEventMsg{agent: row, ev: ev}:
    default:
        s.mu.Lock()
        s.dropped++

Mcmd/entire/cli/review/tui_sink.go+34/-1

536 unmodified lines

537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587

536 unmodified lines

t.Errorf("codex row = %+v, want single {codex,200}", got.AgentRuns[1])
    }
}

// TestTUISink_SumsLiveTokensAcrossWorkers pins that live token counts on a
// collapsed agent row are the SUM of its workers' cumulative counts, not
// last-writer-wins. Since #1666 streams cumulative tokens per worker and the
// row overwrites on each Tokens event, two skill-workers sharing one agent
// row would otherwise bounce between their individual counts mid-run.
func TestTUISink_SumsLiveTokensAcrossWorkers(t *testing.T) {
    t.Parallel()
    prog := newRecordingProgram()
    sink := newTUISinkWithProgram(prog)
    sink.groupWorkers([]string{tAgentClaude}, map[string]string{
        "claude-code:review":    tAgentClaude,
        "claude-code:pr-review": tAgentClaude,
    })
    sink.Start()
    defer func() { prog.Kill(); sink.Wait() }()

sink.AgentEvent("claude-code:review", reviewtypes.Tokens{In: 100, Out: 1})
    sink.AgentEvent("claude-code:pr-review", reviewtypes.Tokens{In: 50, Out: 2})
    sink.AgentEvent("claude-code:review", reviewtypes.Tokens{In: 120, Out: 3}) // worker 1 advances

// Last forwarded Tokens for the agent row must be the running sum:
    // worker1's latest {120,3} + worker2's latest {50,2} = {170,5}.
    want := reviewtypes.Tokens{In: 170, Out: 5}
    deadline := time.After(5 * time.Second)
    for {
        var last *reviewtypes.Tokens
        for _, m := range prog.recorded() {
            if ae, ok := m.(agentEventMsg); ok {
                if tk, ok := ae.ev.(reviewtypes.Tokens); ok {
                    if ae.agent != tAgentClaude {
                        t.Fatalf("tokens routed to %q, want %q", ae.agent, tAgentClaude)
                    }
                    tk := tk
                    last = &tk
                }
            }
        }
        if last != nil && *last == want {
            return
        }
        select {
        case <-deadline:
            t.Fatalf("last forwarded tokens = %v, want %v (per-worker sum)", last, want)
        case <-time.After(10 * time.Millisecond):
        }
    }
}

Mcmd/entire/cli/review/tui_sink_test.go+48