diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index 23040984..8da8fd1e 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -69,9 +69,6 @@ func (r *UsageReporter) publishWithOutcome(ctx context.Context, detail usage.Det detail.TotalTokens = total } } - if detail.InputTokens == 0 && detail.OutputTokens == 0 && detail.ReasoningTokens == 0 && detail.CachedTokens == 0 && detail.TotalTokens == 0 && !failed { - return - } r.once.Do(func() { usage.PublishRecord(ctx, r.buildRecord(detail, failed)) }) diff --git a/test/usage_logging_test.go b/test/usage_logging_test.go new file mode 100644 index 00000000..41c2ee34 --- /dev/null +++ b/test/usage_logging_test.go @@ -0,0 +1,97 @@ +package test + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v6/internal/config" + runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor" + internalusage "github.com/router-for-me/CLIProxyAPI/v6/internal/usage" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator" +) + +func TestGeminiExecutorRecordsSuccessfulZeroUsageInStatistics(t *testing.T) { + model := fmt.Sprintf("gemini-2.5-flash-zero-usage-%d", time.Now().UnixNano()) + source := fmt.Sprintf("zero-usage-%d@example.com", time.Now().UnixNano()) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + wantPath := "/v1beta/models/" + model + ":generateContent" + if r.URL.Path != wantPath { + t.Fatalf("path = %q, want %q", r.URL.Path, wantPath) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":0,"candidatesTokenCount":0,"totalTokenCount":0}}`)) + })) + defer server.Close() + + executor := runtimeexecutor.NewGeminiExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "gemini", + Attributes: map[string]string{ + "api_key": "test-upstream-key", + "base_url": server.URL, + }, + Metadata: map[string]any{ + "email": source, + }, + } + + prevStatsEnabled := internalusage.StatisticsEnabled() + internalusage.SetStatisticsEnabled(true) + t.Cleanup(func() { + internalusage.SetStatisticsEnabled(prevStatsEnabled) + }) + + _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: model, + Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatGemini, + OriginalRequest: []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`), + }) + if err != nil { + t.Fatalf("Execute error: %v", err) + } + + detail := waitForStatisticsDetail(t, "gemini", model, source) + if detail.Failed { + t.Fatalf("detail failed = true, want false") + } + if detail.Tokens.TotalTokens != 0 { + t.Fatalf("total tokens = %d, want 0", detail.Tokens.TotalTokens) + } +} + +func waitForStatisticsDetail(t *testing.T, apiName, model, source string) internalusage.RequestDetail { + t.Helper() + + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + snapshot := internalusage.GetRequestStatistics().Snapshot() + apiSnapshot, ok := snapshot.APIs[apiName] + if !ok { + time.Sleep(10 * time.Millisecond) + continue + } + modelSnapshot, ok := apiSnapshot.Models[model] + if !ok { + time.Sleep(10 * time.Millisecond) + continue + } + for _, detail := range modelSnapshot.Details { + if detail.Source == source { + return detail + } + } + time.Sleep(10 * time.Millisecond) + } + + t.Fatalf("timed out waiting for statistics detail for api=%q model=%q source=%q", apiName, model, source) + return internalusage.RequestDetail{} +}