- 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 | func (r *UsageReporter) PublishAdditionalModel(ctx context.Context, model string, detail usage.Detail) { | 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 | return | 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 | model = strings.TrimSpace(model) | 63 | model = strings.TrimSpace(model) |
| 56 | if model == "" { | 64 | if model == "" { |
| 57 | return | 65 | return usage.Record{}, false |
| 58 | } | 66 | } |
| 59 | detail = normalizeUsageDetailTotal(detail) | 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 | func (r *UsageReporter) PublishFailure(ctx context.Context) { | 74 | func (r *UsageReporter) PublishFailure(ctx context.Context) { |
| @@ -93,6 +104,14 @@ func normalizeUsageDetailTotal(detail usage.Detail) usage.Detail { · 29 unmodified lines | |||
| 93 | return detail | 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 | // ensurePublished guarantees that a usage record is emitted exactly once. | 115 | // ensurePublished guarantees that a usage record is emitted exactly once. |
| 97 | // It is safe to call multiple times; only the first call wins due to once.Do. | 116 | // It is safe to call multiple times; only the first call wins due to once.Do. |
| 98 | // This is used to ensure request counting even when upstream responses do not | 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 | t.Fatalf("latency = %v, want <= 3s", record.Latency) | 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 | } | ||