diff --git a/aibridge/intercept/chatcompletions/base.go b/aibridge/intercept/chatcompletions/base.go index e2c4a82ad4..4e6db3abe6 100644 --- a/aibridge/intercept/chatcompletions/base.go +++ b/aibridge/intercept/chatcompletions/base.go @@ -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. */) } diff --git a/aibridge/intercept/chatcompletions/base_internal_test.go b/aibridge/intercept/chatcompletions/base_internal_test.go index af31a59bdc..55baa15a2a 100644 --- a/aibridge/intercept/chatcompletions/base_internal_test.go +++ b/aibridge/intercept/chatcompletions/base_internal_test.go @@ -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() diff --git a/aibridge/intercept/chatcompletions/blocking.go b/aibridge/intercept/chatcompletions/blocking.go index b3d829377b..d5913557d0 100644 --- a/aibridge/intercept/chatcompletions/blocking.go +++ b/aibridge/intercept/chatcompletions/blocking.go @@ -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 diff --git a/aibridge/intercept/chatcompletions/streaming.go b/aibridge/intercept/chatcompletions/streaming.go index bf76bdcfb4..17ec61edbf 100644 --- a/aibridge/intercept/chatcompletions/streaming.go +++ b/aibridge/intercept/chatcompletions/streaming.go @@ -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 { diff --git a/aibridge/intercept/responses/base.go b/aibridge/intercept/responses/base.go index 0154457795..992e921e90 100644 --- a/aibridge/intercept/responses/base.go +++ b/aibridge/intercept/responses/base.go @@ -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(), diff --git a/aibridge/intercept/responses/base_internal_test.go b/aibridge/intercept/responses/base_internal_test.go index 208ee0afc3..265176d536 100644 --- a/aibridge/intercept/responses/base_internal_test.go +++ b/aibridge/intercept/responses/base_internal_test.go @@ -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 {