fix: skip Claude prompt estimate for zero-frame aborts

relay/channel/claude/relay-claude.go

@@ -150,11 +150,29 @@ func countClaudeStreamBillableTools(c *gin.Context, info *relaycommon.RelayInfo,
 	}
 }
 
+func hasClaudeStreamUsageEvidence(usage *dto.Usage) bool {
+	if usage == nil {
+		return false
+	}
+	return usage.PromptTokens != 0 ||
+		usage.CompletionTokens != 0 ||
+		usage.PromptTokensDetails.CachedTokens != 0 ||
+		usage.PromptTokensDetails.CachedCreationTokens != 0 ||
+		usage.ClaudeCacheCreation5mTokens != 0 ||
+		usage.ClaudeCacheCreation1hTokens != 0 ||
+		usage.BillingUsage != nil
+}
+
 func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo) {
 	if claudeInfo.Usage.PromptTokens == 0 {
 		//上游出错
 	}
-	if claudeInfo.Usage.CompletionTokens == 0 || !claudeInfo.Done {
+	abnormalZeroFrameWithoutUsage := info != nil &&
+		info.StreamStatus != nil &&
+		!info.StreamStatus.IsNormalEnd() &&
+		info.ReceivedResponseCount == 0 &&
+		!hasClaudeStreamUsageEvidence(claudeInfo.Usage)
+	if !abnormalZeroFrameWithoutUsage && (claudeInfo.Usage.CompletionTokens == 0 || !claudeInfo.Done) {
 		if common.DebugEnabled {
 			common.SysLog("claude response usage is not complete, maybe upstream error")
 		}

relay/channel/claude/zero_frame_billing_test.go

@@ -0,0 +1,174 @@
+package claude
+
+import (
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/QuantumNous/new-api/relaykit/dto"
+	"github.com/QuantumNous/new-api/relaykit/types"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+)
+
+func newClaudeFinalResponseTestContext() *gin.Context {
+	c, _ := gin.CreateTestContext(httptest.NewRecorder())
+	return c
+}
+
+func newClaudeFinalResponseTestInfo(receivedResponseCount, estimatePromptTokens int) *relaycommon.RelayInfo {
+	info := &relaycommon.RelayInfo{
+		ReceivedResponseCount: receivedResponseCount,
+		RelayFormat:           types.RelayFormatClaude,
+		ChannelMeta: &relaycommon.ChannelMeta{
+			UpstreamModelName: "claude-test",
+		},
+	}
+	info.SetEstimatePromptTokens(estimatePromptTokens)
+	return info
+}
+
+func TestHandleStreamFinalResponse_ZeroFramesWithoutUsageEvidenceKeepsZeroUsage(t *testing.T) {
+	c := newClaudeFinalResponseTestContext()
+	info := newClaudeFinalResponseTestInfo(0, 1234)
+	info.StreamStatus = relaycommon.NewStreamStatus()
+	info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, nil)
+	claudeInfo := &ClaudeResponseInfo{
+		ResponseText: strings.Builder{},
+		Usage:        &dto.Usage{},
+	}
+	claudeInfo.ResponseText.WriteString("partial response text")
+
+	HandleStreamFinalResponse(c, info, claudeInfo)
+
+	require.Zero(t, claudeInfo.Usage.PromptTokens)
+	require.Zero(t, claudeInfo.Usage.CompletionTokens)
+	require.Zero(t, claudeInfo.Usage.TotalTokens)
+	require.Nil(t, claudeInfo.Usage.BillingUsage)
+	require.False(t, common.GetContextKeyBool(c, constant.ContextKeyLocalCountTokens))
+}
+
+func TestHandleStreamFinalResponse_ZeroCountWithMessageStartUsagePreservesUpstreamEvidence(t *testing.T) {
+	c := newClaudeFinalResponseTestContext()
+	info := newClaudeFinalResponseTestInfo(0, 1234)
+	info.StreamStatus = relaycommon.NewStreamStatus()
+	info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, nil)
+	claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{}}
+
+	require.True(t, FormatClaudeResponseInfo(&dto.ClaudeResponse{
+		Type: "message_start",
+		Message: &dto.ClaudeMediaMessage{
+			Usage: &dto.ClaudeUsage{
+				InputTokens:              100,
+				CacheReadInputTokens:     30,
+				CacheCreationInputTokens: 50,
+			},
+		},
+	}, nil, claudeInfo))
+
+	HandleStreamFinalResponse(c, info, claudeInfo)
+
+	require.Equal(t, 100, claudeInfo.Usage.PromptTokens)
+	require.Equal(t, 30, claudeInfo.Usage.PromptTokensDetails.CachedTokens)
+	require.Equal(t, 50, claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens)
+	require.NotNil(t, claudeInfo.Usage.BillingUsage)
+}
+
+func TestHandleStreamFinalResponse_ZeroFramesWithCacheOnlyUsagePreservesUpstreamEvidence(t *testing.T) {
+	c := newClaudeFinalResponseTestContext()
+	info := newClaudeFinalResponseTestInfo(0, 1234)
+	info.StreamStatus = relaycommon.NewStreamStatus()
+	info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, nil)
+	claudeInfo := &ClaudeResponseInfo{
+		Usage: &dto.Usage{
+			PromptTokensDetails: dto.InputTokenDetails{
+				CachedTokens: 30,
+			},
+		},
+	}
+
+	HandleStreamFinalResponse(c, info, claudeInfo)
+
+	require.Equal(t, 30, claudeInfo.Usage.PromptTokensDetails.CachedTokens)
+	require.NotNil(t, claudeInfo.Usage.BillingUsage)
+}
+
+func TestHandleStreamFinalResponse_ZeroCountWithoutAbnormalStatusKeepsLocalFallback(t *testing.T) {
+	tests := []struct {
+		name      string
+		endReason *relaycommon.StreamEndReason
+	}{
+		{name: "AWS stream status absent"},
+		{name: "normal EOF", endReason: commonPointer(relaycommon.StreamEndReasonEOF)},
+		{name: "normal done", endReason: commonPointer(relaycommon.StreamEndReasonDone)},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := newClaudeFinalResponseTestContext()
+			info := newClaudeFinalResponseTestInfo(0, 1234)
+			if tt.endReason != nil {
+				info.StreamStatus = relaycommon.NewStreamStatus()
+				info.StreamStatus.SetEndReason(*tt.endReason, nil)
+			}
+			claudeInfo := &ClaudeResponseInfo{
+				ResponseText: strings.Builder{},
+				Usage:        &dto.Usage{},
+			}
+			claudeInfo.ResponseText.WriteString("response text")
+
+			HandleStreamFinalResponse(c, info, claudeInfo)
+
+			require.Equal(t, 1234, claudeInfo.Usage.PromptTokens)
+			require.Positive(t, claudeInfo.Usage.CompletionTokens)
+			require.True(t, common.GetContextKeyBool(c, constant.ContextKeyLocalCountTokens))
+		})
+	}
+}
+
+func TestHandleStreamFinalResponse_ReceivedFramesWithoutUsageKeepsLocalFallback(t *testing.T) {
+	c := newClaudeFinalResponseTestContext()
+	info := newClaudeFinalResponseTestInfo(1, 1234)
+	info.StreamStatus = relaycommon.NewStreamStatus()
+	info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, nil)
+	claudeInfo := &ClaudeResponseInfo{
+		ResponseText: strings.Builder{},
+		Usage:        &dto.Usage{},
+	}
+	claudeInfo.ResponseText.WriteString("response text")
+
+	HandleStreamFinalResponse(c, info, claudeInfo)
+
+	require.Equal(t, 1234, claudeInfo.Usage.PromptTokens)
+	require.Positive(t, claudeInfo.Usage.CompletionTokens)
+	require.True(t, common.GetContextKeyBool(c, constant.ContextKeyLocalCountTokens))
+}
+
+func TestHandleStreamFinalResponse_CompleteUpstreamUsageRemainsUnchanged(t *testing.T) {
+	c := newClaudeFinalResponseTestContext()
+	info := newClaudeFinalResponseTestInfo(3, 1234)
+	billingUsage := dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{
+		InputTokens:  100,
+		OutputTokens: 20,
+	})
+	claudeInfo := &ClaudeResponseInfo{
+		Done: true,
+		Usage: &dto.Usage{
+			PromptTokens:     100,
+			CompletionTokens: 20,
+			TotalTokens:      120,
+			BillingUsage:     billingUsage,
+		},
+	}
+
+	HandleStreamFinalResponse(c, info, claudeInfo)
+
+	require.Equal(t, 100, claudeInfo.Usage.PromptTokens)
+	require.Equal(t, 20, claudeInfo.Usage.CompletionTokens)
+	require.Equal(t, 120, claudeInfo.Usage.TotalTokens)
+	require.Same(t, billingUsage, claudeInfo.Usage.BillingUsage)
+	require.False(t, common.GetContextKeyBool(c, constant.ContextKeyLocalCountTokens))
+}