Deduplicate no-id Pi review token events

40cb4d6→main·

dipree·2w ago·2 files·+70 added/-7 removed

Sessions

90ce539be7b6View transcript

[?
Trail Finding and Pi Review Token DeduplicationPi·GPT-5.5·1 step](/content/gh/entireio/cli/session/019f1893-b810-7992-afb4-8c4bddc4ae3c#timeline-90ce539be7b6/index.html)

Changes

2

• cmd/entire/cli/agent/pi

• Mreviewer.go+45/-7

• Mreviewer_test.go+25

50 unmodified lines

51
52
53
54
55
56
57
58
10 unmodified lines

69
70
71
70
72
73
74
75
76
77
6 unmodified lines

84
85
86
83
87
88
89
90
12 unmodified lines

103
104
105
102
106
107
108
109
66 unmodified lines

176
177
178
175
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
181
211
212
213
214
215
185
186
216
217
218
219
220
221
222
223
224
225
226
227

50 unmodified lines

scanner.Buffer(make([]byte, min(1024*1024, piReviewMaxScannerBuf)), piReviewMaxScannerBuf)
        messageIDsWithTextDelta := map[string]struct{}{}
        messageIDsWithUsage := map[string]struct{}{}
        messageUsageByTurn := map[int]map[piReviewUsageKey]struct{}{}
        urnNumber := 0
        tokens := reviewtypes.Tokens{}
        finished := false
        success := true
10 unmodified lines

}

switch env.Type {
        case "session", "agent_start", "turn_start", "queue_update", "compaction_start", "compaction_end", "auto_retry_start", "auto_retry_end":
        case "turn_start":
            urnNumber++
        case "session", "agent_start", "queue_update", "compaction_start", "compaction_end", "auto_retry_start", "auto_retry_end":
            // Session/control events do not map to user-visible review output.
        case "message_update":
            if text := env.AssistantMessageEvent.TextDelta(); text != "" {
6 unmodified lines

success = false
            }
            if env.Message.Usage != nil {
                emitPiReviewTokens(out, env, &tokens, messageIDsWithUsage)
                emitPiReviewTokens(out, env, &tokens, messageIDsWithUsage, messageUsageByTurn, turnNumber)
            }
            if _, sawDelta := messageIDsWithTextDelta[env.MessageID()]; !sawDelta {
                if text := piReviewMessageText(env.Message.Content); text != "" {
12 unmodified lines

success = false
                }
                if env.Message.Usage != nil {
                    emitPiReviewTokens(out, env, &tokens, messageIDsWithUsage)
                    emitPiReviewTokens(out, env, &tokens, messageIDsWithUsage, messageUsageByTurn, turnNumber)
                }
            case "agent_end":
                finished = true
66 unmodified lines

CacheWrite int `json:"cacheWrite"`
}

func emitPiReviewTokens(out chan<- reviewtypes.Event, env piReviewEnvelope, total *reviewtypes.Tokens, seen map[string]struct{}) {
type piReviewUsageKey struct {
    Input      int
    Output     int
    CacheRead  int
    CacheWrite int
}

func emitPiReviewTokens(out chan<- reviewtypes.Event, env piReviewEnvelope, total *reviewtypes.Tokens, seen map[string]struct{}, messageUsageByTurn map[int]map[piReviewUsageKey]struct{}, turnNumber int) {
    if env.Message.Usage == nil || total == nil {
        return
    }
    if shouldSkipPiReviewUsage(env, seen, messageUsageByTurn, turnNumber) {
        return
    }
    *total = addPiReviewTokens(*total, env.Message.Usage)
    out <- *total
}

func shouldSkipPiReviewUsage(env piReviewEnvelope, seen map[string]struct{}, messageUsageByTurn map[int]map[piReviewUsageKey]struct{}, turnNumber int) bool {
    usage := env.Message.Usage
    if usage == nil {
        return true
    }
    sig := piReviewUsageKey{Input: usage.Input, Output: usage.Output, CacheRead: usage.CacheRead, CacheWrite: usage.CacheWrite}
    if env.Type == "message_end" && messageUsageByTurn != nil {
        if messageUsageByTurn[turnNumber] == nil {
            messageUsageByTurn[turnNumber] = map[piReviewUsageKey]struct{}{}
        }
        messageUsageByTurn[turnNumber][sig] = struct{}{}
    }
    if key := env.MessageID(); key != "" {
        if _, ok := seen[key]; ok {
            return
            return true
        }
        seen[key] = struct{}{}
        return false
    }
    *total = addPiReviewTokens(*total, env.Message.Usage)
    out <- *total
    // Pi streams can emit usage on both message_end and turn_end. Some realistic
    // streams omit ids on both events, so fall back to the current turn's usage
    // signature to avoid counting a no-id turn_end duplicate of the message_end.
    if env.Type == "turn_end" && messageUsageByTurn != nil {
        if _, ok := messageUsageByTurn[turnNumber][sig]; ok {
            return true
        }
    }
    return false
}

func addPiReviewTokens(total reviewtypes.Tokens, usage *piReviewUsage) reviewtypes.Tokens {

Mcmd/entire/cli/agent/pi/reviewer.go+45/-7

168 unmodified lines

169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199

168 unmodified lines

}
}

func TestPiReviewer_ParseTokensDedupesNoIDTurnEndForSameUsage(t *testing.T) {
    t.Parallel()
    input := strings.Join([]string{
        `{"type":"agent_start"}`,
        `{"type":"turn_start"}`,
        `{"type":"message_end","message":{"role":"assistant","usage":{"input":10,"output":5,"cacheRead":2,"cacheWrite":1},"stopReason":"stop"}}`, 
        `{"type":"turn_end","message":{"role":"assistant","usage":{"input":10,"output":5,"cacheRead":2,"cacheWrite":1},"stopReason":"stop"}}`,
        `{"type":"agent_end"}`,
    }, "\n")

events := collectPiReviewEvents(input)
    var tokens []reviewtypes.Tokens
    for _, ev := range events {
        if tok, ok := ev.(reviewtypes.Tokens); ok {
            tokens = append(tokens, tok)
        }
    }
    if len(tokens) != 1 {
        t.Fatalf("token events = %d, want 1: %#v", len(tokens), events)
    }
    if got := tokens[0]; got.In != 10 || got.Out != 5 {
        t.Fatalf("Tokens = %#v, want In=10 Out=5", got)
    }
}

func collectPiReviewEvents(input string) []reviewtypes.Event {
    ch := parsePiReviewOutput(strings.NewReader(input))
    var events []reviewtypes.Event