- Added `buildAdditionalModelRecord` to filter out zero-token usage details. - Introduced `hasNonZeroTokenUsage` helper function for token usage validation. - Updated tests to cover scenarios for zero and non-zero token usage.
Unified · Split
| @@ -49,15 +49,26 @@ func (r *UsageReporter) Publish(ctx context.Context, detail usage.Detail) { · 48 unmodified lines | ||
| 49 | 49 | } |
| 50 | 50 | |
| 51 | 51 | func (r *UsageReporter) PublishAdditionalModel(ctx context.Context, model string, detail usage.Detail) { |
| 52 | − if r == nil { | |
| 52 | + record, ok := r.buildAdditionalModelRecord(model, detail) | |
| 53 | + if !ok { | |
| 53 | 54 | return |
| 54 | 55 | } |
| 56 | + usage.PublishRecord(ctx, record) | |
| 57 | +} | |
| 58 | + | |
| 59 | +func (r *UsageReporter) buildAdditionalModelRecord(model string, detail usage.Detail) (usage.Record, bool) { | |
| 60 | + if r == nil { | |
| 61 | + return usage.Record{}, false | |
| 62 | + } | |
| 55 | 63 | model = strings.TrimSpace(model) |
| 56 | 64 | if model == "" { |
| 57 | − return | |
| 65 | + return usage.Record{}, false | |
| 58 | 66 | } |
| 59 | 67 | detail = normalizeUsageDetailTotal(detail) |
| 60 | − usage.PublishRecord(ctx, r.buildRecordForModel(model, detail, false)) | |
| 68 | + if !hasNonZeroTokenUsage(detail) { | |
| 69 | + return usage.Record{}, false | |
| 70 | + } | |
| 71 | + return r.buildRecordForModel(model, detail, false), true | |
| 61 | 72 | } |
| 62 | 73 | |
| 63 | 74 | func (r *UsageReporter) PublishFailure(ctx context.Context) { |
| @@ -93,6 +104,14 @@ func normalizeUsageDetailTotal(detail usage.Detail) usage.Detail { · 29 unmodified lines | ||
| 93 | 104 | return detail |
| 94 | 105 | } |
| 95 | 106 | |
| 107 | +func hasNonZeroTokenUsage(detail usage.Detail) bool { | |
| 108 | + return detail.InputTokens != 0 || | |
| 109 | + detail.OutputTokens != 0 || | |
| 110 | + detail.ReasoningTokens != 0 || | |
| 111 | + detail.CachedTokens != 0 || | |
| 112 | + detail.TotalTokens != 0 | |
| 113 | +} | |
| 114 | + | |
| 96 | 115 | // ensurePublished guarantees that a usage record is emitted exactly once. |
| 97 | 116 | // It is safe to call multiple times; only the first call wins due to once.Do. |
| 98 | 117 | // This is used to ensure request counting even when upstream responses do not |
| @@ -62,3 +62,21 @@ func TestUsageReporterBuildRecordIncludesLatency(t *testing.T) { · 61 unmodified lines | ||
| 62 | 62 | t.Fatalf("latency = %v, want <= 3s", record.Latency) |
| 63 | 63 | } |
| 64 | 64 | } |
| 65 | + | |
| 66 | +func TestUsageReporterBuildAdditionalModelRecordSkipsZeroTokens(t *testing.T) { | |
| 67 | + reporter := &UsageReporter{ | |
| 68 | + provider: "codex", | |
| 69 | + model: "gpt-5.4", | |
| 70 | + requestedAt: time.Now(), | |
| 71 | + } | |
| 72 | + | |
| 73 | + if _, ok := reporter.buildAdditionalModelRecord("gpt-image-2", usage.Detail{}); ok { | |
| 74 | + t.Fatalf("expected all-zero token usage to be skipped") | |
| 75 | + } | |
| 76 | + if _, ok := reporter.buildAdditionalModelRecord("gpt-image-2", usage.Detail{InputTokens: 2}); !ok { | |
| 77 | + t.Fatalf("expected non-zero input token usage to be recorded") | |
| 78 | + } | |
| 79 | + if _, ok := reporter.buildAdditionalModelRecord("gpt-image-2", usage.Detail{CachedTokens: 2}); !ok { | |
| 80 | + t.Fatalf("expected non-zero cached token usage to be recorded") | |
| 81 | + } | |
| 82 | +} | |