feat: support cached input token billing

2026-06-23 17:17:57 +08:00
parent 6089aa6085
commit 7f32446466
16 changed files with 570 additions and 45 deletions
@@ -165,6 +165,87 @@ func TestWriteCompatibleTaskResponseReturnsSSEWhenStreamIsTrue(t *testing.T) {
 	}
 }

+func TestWriteCompatibleTaskResponseStreamsCachedUsageProviderShapes(t *testing.T) {
+	cases := []struct {
+		name     string
+		usage    map[string]any
+		fragment string
+	}{
+		{
+			name: "aliyun dashscope",
+			usage: map[string]any{
+				"prompt_tokens":     10,
+				"completion_tokens": 2,
+				"total_tokens":      12,
+				"prompt_tokens_details": map[string]any{
+					"cached_tokens": 7,
+				},
+			},
+			fragment: `"prompt_tokens_details":{"cached_tokens":7}`,
+		},
+		{
+			name: "volces responses",
+			usage: map[string]any{
+				"prompt_tokens":     9,
+				"completion_tokens": 3,
+				"total_tokens":      12,
+				"input_tokens_details": map[string]any{
+					"cached_tokens": 6,
+				},
+			},
+			fragment: `"input_tokens_details":{"cached_tokens":6}`,
+		},
+		{
+			name: "deepseek",
+			usage: map[string]any{
+				"prompt_tokens":            8,
+				"completion_tokens":        4,
+				"total_tokens":             12,
+				"prompt_cache_hit_tokens":  5,
+				"prompt_cache_miss_tokens": 3,
+			},
+			fragment: `"prompt_cache_hit_tokens":5`,
+		},
+		{
+			name: "glm",
+			usage: map[string]any{
+				"prompt_tokens":     11,
+				"completion_tokens": 1,
+				"total_tokens":      12,
+				"prompt_tokens_details": map[string]any{
+					"cached_tokens": 4,
+				},
+			},
+			fragment: `"prompt_tokens_details":{"cached_tokens":4}`,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			executor := &fakeTaskExecutor{
+				deltas: []clients.StreamDeltaEvent{{Text: "ok"}},
+				output: map[string]any{
+					"id":     "chatcmpl-cache",
+					"object": "chat.completion",
+					"usage":  tc.usage,
+				},
+			}
+			req := httptest.NewRequest(http.MethodPost, "/api/v1/chat/completions", nil)
+			recorder := httptest.NewRecorder()
+
+			writeCompatibleTaskResponse(context.Background(), recorder, req, executor, "chat.completions", "cache-test", store.GatewayTask{ID: "task-test"}, &auth.User{}, true, true)
+
+			body := recorder.Body.String()
+			if !strings.Contains(body, tc.fragment) {
+				t.Fatalf("SSE body missing cached usage fragment %s: %s", tc.fragment, body)
+			}
+			if !strings.Contains(body, `"choices":[]`) || !strings.Contains(body, "data: [DONE]") {
+				t.Fatalf("SSE body should include final usage chunk and done marker: %s", body)
+			}
+		})
+	}
+}
+
 func TestWriteCompatibleTaskResponseStreamsStructuredToolAndReasoningDeltas(t *testing.T) {
 	executor := &fakeTaskExecutor{
 		deltas: []clients.StreamDeltaEvent{