opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: negative metric counter increment from token arithmetic (#26547)

Browse Source 
_Disclosure: produced using Claude Opus 4.8_

Closes
[AIGOV-452](https://linear.app/codercom/issue/AIGOV-452/prevent-control-plane-panic-on-negative-cached-tokens)

Also addresses a similar shortcoming in
`aibridge/intercept/responses/base.go` and aligns token recording
approach for chatcompletions with other implementations

---------

Signed-off-by: Danny Kopping <danny@coder.com>
This commit is contained in:
Danny Kopping
2026-06-19 15:44:04 +02:00
committed by GitHub
parent a12b051834
commit 6186532cec
6 changed files with 177 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
aibridge/intercept/chatcompletions/base.go
+21 -2
View File
@@ -226,6 +226,25 @@ func (i *interceptionBase) hasInjectableTools() bool {
return i.mcpProxy != nil && len(i.mcpProxy.ListTools()) > 0
}
// recordTokenUsage records the token usage for a single completion, accounting
// for cached tokens included in the prompt token count.
func (i *interceptionBase) recordTokenUsage(ctx context.Context, msgID string, usage openai.CompletionUsage) {
_ = i.recorder.RecordTokenUsage(ctx, &recorder.TokenUsageRecord{
InterceptionID: i.ID().String(),
MsgID: msgID,
Input: calculateActualInputTokenUsage(usage),
Output: usage.CompletionTokens,
CacheReadInputTokens: usage.PromptTokensDetails.CachedTokens,
ExtraTokenTypes: map[string]int64{
"prompt_audio": usage.PromptTokensDetails.AudioTokens,
"completion_accepted_prediction": usage.CompletionTokensDetails.AcceptedPredictionTokens,
"completion_rejected_prediction": usage.CompletionTokensDetails.RejectedPredictionTokens,
"completion_audio": usage.CompletionTokensDetails.AudioTokens,
"completion_reasoning": usage.CompletionTokensDetails.ReasoningTokens,
},
})
}
func sumUsage(ref, in openai.CompletionUsage) openai.CompletionUsage {
return openai.CompletionUsage{
CompletionTokens: ref.CompletionTokens + in.CompletionTokens,
@@ -249,6 +268,6 @@ func calculateActualInputTokenUsage(in openai.CompletionUsage) int64 {
// Input *includes* the cached tokens, so we subtract them here to reflect actual input token usage.
// The original value can be reconstructed by adding CachedTokens back to Input.
// See https://platform.openai.com/docs/api-reference/usage/completions_object#usage/completions_object-input_tokens.
return in.PromptTokens /* The aggregated number of text input tokens used, including cached tokens. */ -
in.PromptTokensDetails.CachedTokens /* The aggregated number of text input tokens that has been cached from previous requests. */
return max(0, in.PromptTokens /* The aggregated number of text input tokens used, including cached tokens. */ -
in.PromptTokensDetails.CachedTokens /* The aggregated number of text input tokens that has been cached from previous requests. */)
}
aibridge/intercept/chatcompletions/base_internal_test.go
+125
View File
@@ -7,6 +7,7 @@ import (
"testing"
"time"
"github.com/google/uuid"
"github.com/openai/openai-go/v3"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -15,11 +16,135 @@ import (
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/aibridge/config"
"github.com/coder/coder/v2/aibridge/intercept"
"github.com/coder/coder/v2/aibridge/internal/testutil"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/recorder"
"github.com/coder/coder/v2/aibridge/utils"
"github.com/coder/quartz"
)
func TestRecordTokenUsage(t *testing.T) {
t.Parallel()
id := uuid.MustParse("22222222-2222-2222-2222-222222222222")
tests := []struct {
name string
msgID string
usage openai.CompletionUsage
expected *recorder.TokenUsageRecord
}{
{
name: "with_all_token_details",
msgID: "cmpl_full",
usage: openai.CompletionUsage{
PromptTokens: 100,
CompletionTokens: 50,
TotalTokens: 150,
PromptTokensDetails: openai.CompletionUsagePromptTokensDetails{
CachedTokens: 40,
AudioTokens: 3,
},
CompletionTokensDetails: openai.CompletionUsageCompletionTokensDetails{
AcceptedPredictionTokens: 7,
RejectedPredictionTokens: 2,
AudioTokens: 1,
ReasoningTokens: 9,
},
},
expected: &recorder.TokenUsageRecord{
InterceptionID: id.String(),
MsgID: "cmpl_full",
Input: 60, // 100 prompt - 40 cached
Output: 50,
CacheReadInputTokens: 40,
ExtraTokenTypes: map[string]int64{
"prompt_audio": 3,
"completion_accepted_prediction": 7,
"completion_rejected_prediction": 2,
"completion_audio": 1,
"completion_reasoning": 9,
},
},
},
{
name: "all_tokens_cached",
msgID: "cmpl_cached",
usage: openai.CompletionUsage{
PromptTokens: 100,
CompletionTokens: 20,
PromptTokensDetails: openai.CompletionUsagePromptTokensDetails{
CachedTokens: 100,
},
},
expected: &recorder.TokenUsageRecord{
InterceptionID: id.String(),
MsgID: "cmpl_cached",
Input: 0, // 100 prompt - 100 cached
Output: 20,
CacheReadInputTokens: 100,
ExtraTokenTypes: map[string]int64{
"prompt_audio": 0,
"completion_accepted_prediction": 0,
"completion_rejected_prediction": 0,
"completion_audio": 0,
"completion_reasoning": 0,
},
},
},
{
// Upstream violates the invariant that PromptTokens includes
// CachedTokens. Input must clamp to 0 so it never panics a
// Prometheus counter when used as an increment.
name: "cached_tokens_exceed_prompt_tokens_clamps_to_zero",
msgID: "cmpl_clamp",
usage: openai.CompletionUsage{
PromptTokens: 40,
CompletionTokens: 20,
PromptTokensDetails: openai.CompletionUsagePromptTokensDetails{
CachedTokens: 100,
},
},
expected: &recorder.TokenUsageRecord{
InterceptionID: id.String(),
MsgID: "cmpl_clamp",
Input: 0, // max(0, 40 prompt - 100 cached)
Output: 20,
CacheReadInputTokens: 100,
ExtraTokenTypes: map[string]int64{
"prompt_audio": 0,
"completion_accepted_prediction": 0,
"completion_rejected_prediction": 0,
"completion_audio": 0,
"completion_reasoning": 0,
},
},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
rec := &testutil.MockRecorder{}
base := &interceptionBase{
id: id,
recorder: rec,
logger: slog.Make(),
}
base.recordTokenUsage(t.Context(), tc.msgID, tc.usage)
tokens := rec.RecordedTokenUsages()
require.Len(t, tokens, 1)
got := tokens[0]
got.CreatedAt = time.Time{} // ignore time
require.Equal(t, tc.expected, got)
require.GreaterOrEqual(t, got.Input, int64(0), "input must never be negative")
})
}
}
func TestScanForCorrelatingToolCallID(t *testing.T) {
t.Parallel()
aibridge/intercept/chatcompletions/blocking.go
+1 -14
View File
@@ -123,20 +123,7 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req
lastUsage := completion.Usage
cumulativeUsage = sumUsage(cumulativeUsage, completion.Usage)
_ = i.recorder.RecordTokenUsage(ctx, &recorder.TokenUsageRecord{
InterceptionID: i.ID().String(),
MsgID: completion.ID,
Input: calculateActualInputTokenUsage(lastUsage),
Output: lastUsage.CompletionTokens,
CacheReadInputTokens: lastUsage.PromptTokensDetails.CachedTokens,
ExtraTokenTypes: map[string]int64{
"prompt_audio": lastUsage.PromptTokensDetails.AudioTokens,
"completion_accepted_prediction": lastUsage.CompletionTokensDetails.AcceptedPredictionTokens,
"completion_rejected_prediction": lastUsage.CompletionTokensDetails.RejectedPredictionTokens,
"completion_audio": lastUsage.CompletionTokensDetails.AudioTokens,
"completion_reasoning": lastUsage.CompletionTokensDetails.ReasoningTokens,
},
})
i.recordTokenUsage(ctx, completion.ID, lastUsage)
// Check if we have tool calls to process.
var pendingToolCalls []openai.ChatCompletionMessageToolCallUnion
aibridge/intercept/chatcompletions/streaming.go
+1 -14
View File
@@ -270,20 +270,7 @@ func (i *StreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Re
if lastUsage := processor.getLastUsage(); lastUsage.CompletionTokens > 0 {
// If the usage information is set, track it.
// The API will send usage information when the response terminates, which will happen if a tool call is invoked.
_ = i.recorder.RecordTokenUsage(streamCtx, &recorder.TokenUsageRecord{
InterceptionID: i.ID().String(),
MsgID: processor.getMsgID(),
Input: calculateActualInputTokenUsage(lastUsage),
Output: lastUsage.CompletionTokens,
CacheReadInputTokens: lastUsage.PromptTokensDetails.CachedTokens,
ExtraTokenTypes: map[string]int64{
"prompt_audio": lastUsage.PromptTokensDetails.AudioTokens,
"completion_accepted_prediction": lastUsage.CompletionTokensDetails.AcceptedPredictionTokens,
"completion_rejected_prediction": lastUsage.CompletionTokensDetails.RejectedPredictionTokens,
"completion_audio": lastUsage.CompletionTokensDetails.AudioTokens,
"completion_reasoning": lastUsage.CompletionTokensDetails.ReasoningTokens,
},
})
i.recordTokenUsage(streamCtx, processor.getMsgID(), lastUsage)
}
if iterationStarted {
aibridge/intercept/responses/base.go
+1 -1
View File
@@ -303,7 +303,7 @@ func (i *responsesInterceptionBase) recordTokenUsage(ctx context.Context, respon
// Keeping logic consistent with chat completions
// Input *includes* the cached tokens, so we subtract them here to reflect actual input token usage.
inputNonCacheTokens := usage.InputTokens - usage.InputTokensDetails.CachedTokens
inputNonCacheTokens := max(0, usage.InputTokens-usage.InputTokensDetails.CachedTokens)
if err := i.recorder.RecordTokenUsage(ctx, &recorder.TokenUsageRecord{
InterceptionID: i.ID().String(),
aibridge/intercept/responses/base_internal_test.go
+28
View File
@@ -318,6 +318,34 @@ func TestRecordTokenUsage(t *testing.T) {
},
},
},
{
// Upstream violates the invariant that InputTokens includes
// CachedTokens. Input must clamp to 0 so it never panics a
// Prometheus counter when used as an increment.
name: "cached_tokens_exceed_input_tokens_clamps_to_zero",
response: &oairesponses.Response{
ID: "resp_clamp",
Usage: oairesponses.ResponseUsage{
InputTokens: 10,
OutputTokens: 20,
TotalTokens: 30,
InputTokensDetails: oairesponses.ResponseUsageInputTokensDetails{
CachedTokens: 40,
},
},
},
expected: &recorder.TokenUsageRecord{
InterceptionID: id.String(),
MsgID: "resp_clamp",
Input: 0, // max(0, 10 input - 40 cached)
Output: 20,
CacheReadInputTokens: 40,
ExtraTokenTypes: map[string]int64{
"output_reasoning": 0,
"total_tokens": 30,
},
},
},
}
for _, tc := range tests {