fix: 合并同一请求的用量与生命周期记录

2026-08-15 00:42:50 +08:00
parent f68741f190
commit eec85bfba9
2 changed files with 263 additions and 2 deletions
@@ -90,6 +90,17 @@ CREATE TABLE IF NOT EXISTS model_prices (
 );
 `

+const (
+	// CLIProxyAPI UsageRecord currently has no RequestID. The usage and terminal
+	// callbacks nevertheless carry request-start timestamps from the same host
+	// execution, normally only a few milliseconds apart. This window is used
+	// only for the read projection and never changes persisted billing facts.
+	orphanLifecycleMatchWindow = 250 * time.Millisecond
+	// Near-equal candidates are deliberately left separate instead of risking a
+	// cross-request association under concurrent same-model traffic.
+	orphanLifecycleTieWindow = time.Millisecond
+)
+
 // SQLiteUsageRepository 使用单写连接保存全部用量记录。
 type SQLiteUsageRepository struct {
 	db *sql.DB
@@ -217,6 +228,12 @@ func (r *SQLiteUsageRepository) ListRecent(ctx context.Context, limit int) ([]co
 	if limit < 1 {
 		return []collection.Record{}, nil
 	}
+	// A logical request can occupy two source rows until projection time, so read
+	// ahead before applying the caller's logical-request limit.
+	scanLimit := limit * 2
+	if scanLimit < limit {
+		scanLimit = limit
+	}
 	rows, err := r.db.QueryContext(ctx, `
 WITH combined AS (
    SELECT
@@ -259,13 +276,13 @@ SELECT request_id, execution_id, trace_id, requested_at, api_key, key_alias, mod
       price_multiplier_denominator, client_ip, outcome, status_code, error
 FROM combined
 ORDER BY requested_at DESC
-LIMIT ?`, limit)
+LIMIT ?`, scanLimit)
 	if err != nil {
 		return nil, fmt.Errorf("查询最近用量记录: %w", err)
 	}
 	defer rows.Close()

-	records := make([]collection.Record, 0, limit)
+	records := make([]collection.Record, 0, scanLimit)
 	for rows.Next() {
 		var record collection.Record
 		var requestedAt string
@@ -301,9 +318,126 @@ LIMIT ?`, limit)
 	if err := rows.Err(); err != nil {
 		return nil, fmt.Errorf("遍历用量记录: %w", err)
 	}
+	records = mergeOrphanLifecycleRecords(records)
+	if len(records) > limit {
+		records = records[:limit]
+	}
 	return records, nil
 }

+type lifecycleMatch struct {
+	index     int
+	distance  time.Duration
+	ambiguous bool
+}
+
+func mergeOrphanLifecycleRecords(records []collection.Record) []collection.Record {
+	usageMatches := make(map[int]lifecycleMatch)
+	requestMatches := make(map[int]lifecycleMatch)
+
+	for usageIndex := range records {
+		if !isOrphanUsage(records[usageIndex]) {
+			continue
+		}
+		for requestIndex := range records {
+			if !isLifecycleOnly(records[requestIndex]) || !sameProjectionModel(records[usageIndex], records[requestIndex]) {
+				continue
+			}
+			distance := records[usageIndex].RequestedAt.Sub(records[requestIndex].RequestedAt)
+			if distance < 0 {
+				distance = -distance
+			}
+			if distance > orphanLifecycleMatchWindow {
+				continue
+			}
+			if current, exists := usageMatches[usageIndex]; exists {
+				usageMatches[usageIndex] = betterLifecycleMatch(current, requestIndex, distance)
+			} else {
+				usageMatches[usageIndex] = lifecycleMatch{index: requestIndex, distance: distance}
+			}
+			if current, exists := requestMatches[requestIndex]; exists {
+				requestMatches[requestIndex] = betterLifecycleMatch(current, usageIndex, distance)
+			} else {
+				requestMatches[requestIndex] = lifecycleMatch{index: usageIndex, distance: distance}
+			}
+		}
+	}
+
+	remove := make(map[int]struct{})
+	for usageIndex, requestMatch := range usageMatches {
+		if requestMatch.ambiguous {
+			continue
+		}
+		usageMatch, ok := requestMatches[requestMatch.index]
+		if !ok || usageMatch.ambiguous || usageMatch.index != usageIndex {
+			continue
+		}
+		records[usageIndex] = mergeLifecycleIntoUsage(records[usageIndex], records[requestMatch.index])
+		remove[requestMatch.index] = struct{}{}
+	}
+
+	if len(remove) == 0 {
+		return records
+	}
+	merged := make([]collection.Record, 0, len(records)-len(remove))
+	for index, record := range records {
+		if _, drop := remove[index]; !drop {
+			merged = append(merged, record)
+		}
+	}
+	return merged
+}
+
+func betterLifecycleMatch(current lifecycleMatch, candidate int, distance time.Duration) lifecycleMatch {
+	if distance+orphanLifecycleTieWindow < current.distance {
+		return lifecycleMatch{index: candidate, distance: distance}
+	}
+	if current.distance+orphanLifecycleTieWindow < distance {
+		return current
+	}
+	current.ambiguous = true
+	return current
+}
+
+func isOrphanUsage(record collection.Record) bool {
+	if strings.TrimSpace(record.RequestID) != "" {
+		return false
+	}
+	return record.CostMicros != nil || record.TotalTokens != 0 || record.InputTokens != 0 ||
+		record.OutputTokens != 0 || record.CacheReadTokens != 0 || record.CacheWriteTokens != 0 ||
+		strings.TrimSpace(record.APIKey) != "" || strings.TrimSpace(record.ExecutorType) != ""
+}
+
+func isLifecycleOnly(record collection.Record) bool {
+	return strings.TrimSpace(record.RequestID) != "" && strings.TrimSpace(record.ExecutionID) == "" &&
+		strings.TrimSpace(record.APIKey) == "" && record.CostMicros == nil && record.InputTokens == 0 &&
+		record.OutputTokens == 0 && record.CacheReadTokens == 0 && record.CacheWriteTokens == 0 &&
+		record.TotalTokens == 0 && strings.TrimSpace(record.Outcome) != ""
+}
+
+func sameProjectionModel(left, right collection.Record) bool {
+	leftModel := strings.ToLower(strings.TrimSpace(left.Model))
+	rightModel := strings.ToLower(strings.TrimSpace(right.Model))
+	return leftModel != "" && leftModel == rightModel
+}
+
+func mergeLifecycleIntoUsage(usage, lifecycle collection.Record) collection.Record {
+	usage.RequestID = lifecycle.RequestID
+	if usage.TraceID == "" {
+		usage.TraceID = lifecycle.TraceID
+	}
+	usage.Failed = usage.Failed || lifecycle.Failed
+	usage.RequestType = lifecycle.RequestType
+	usage.Endpoint = lifecycle.Endpoint
+	usage.Outcome = lifecycle.Outcome
+	usage.StatusCode = lifecycle.StatusCode
+	usage.Error = lifecycle.Error
+	if usage.Latency == 0 {
+		usage.Latency = lifecycle.Latency
+	}
+	return usage
+}
+
 func (r *SQLiteUsageRepository) Close() error {
 	return r.db.Close()
 }
@@ -92,3 +92,130 @@ func TestSQLiteUsageRoundTripsNullableCostAndDurations(t *testing.T) {
 		t.Fatalf("round trip mismatch: %+v", got)
 	}
 }
+
+func TestSQLiteUsageMergesOrphanUsageWithLifecycleProjection(t *testing.T) {
+	for _, outcome := range []struct {
+		name       string
+		value      string
+		statusCode int
+		failed     bool
+	}{
+		{name: "succeeded", value: "succeeded", statusCode: 200},
+		{name: "failed", value: "failed", statusCode: 502, failed: true},
+		{name: "canceled", value: "canceled", statusCode: 499, failed: true},
+	} {
+		for _, order := range []string{"usage-first", "lifecycle-first"} {
+			t.Run(outcome.name+"/"+order, func(t *testing.T) {
+				store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
+				if err != nil {
+					t.Fatal(err)
+				}
+				defer store.Close()
+
+				startedAt := time.Date(2026, 8, 14, 15, 26, 20, 593_000_000, time.UTC)
+				cost := int64(12_345)
+				usage := collection.Record{
+					RequestedAt: startedAt.Add(5 * time.Millisecond), APIKey: "000000", Model: "deepseek-v4-flash",
+					ExecutorType: "CodexExecutor", InputTokens: 19_464, OutputTokens: 79, TotalTokens: 19_543,
+					TTFT: 266 * time.Millisecond, CostMicros: &cost,
+				}
+				lifecycle := collection.RequestRecord{
+					RequestID: "request-1", TraceID: "trace-1", RequestedAt: startedAt,
+					CompletedAt: startedAt.Add(2 * time.Second), Model: "deepseek-v4-flash",
+					SourceFormat: "openai-response", Stream: true, Outcome: outcome.value,
+					StatusCode: outcome.statusCode, Endpoint: "/v1/responses",
+				}
+				insertUsage := func() {
+					if err := store.Insert(context.Background(), usage); err != nil {
+						t.Fatal(err)
+					}
+				}
+				insertLifecycle := func() {
+					if err := store.UpsertRequest(context.Background(), lifecycle); err != nil {
+						t.Fatal(err)
+					}
+				}
+				if order == "usage-first" {
+					insertUsage()
+					insertLifecycle()
+				} else {
+					insertLifecycle()
+					insertUsage()
+				}
+
+				records, err := store.ListRecent(context.Background(), 10)
+				if err != nil {
+					t.Fatal(err)
+				}
+				if len(records) != 1 {
+					t.Fatalf("records = %d, want one logical request: %+v", len(records), records)
+				}
+				got := records[0]
+				if got.RequestID != "request-1" || got.TraceID != "trace-1" || got.APIKey != "000000" ||
+					got.RequestType != "SSE" || got.Endpoint != "/v1/responses" || got.Outcome != outcome.value ||
+					got.StatusCode != outcome.statusCode || got.Failed != outcome.failed ||
+					got.TotalTokens != 19_543 || got.CostMicros == nil || *got.CostMicros != cost {
+					t.Fatalf("merged record = %+v", got)
+				}
+			})
+		}
+	}
+}
+
+func TestSQLiteUsageLeavesAmbiguousConcurrentRequestsSeparate(t *testing.T) {
+	store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer store.Close()
+
+	startedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC)
+	for index := 0; index < 2; index++ {
+		if err := store.Insert(context.Background(), collection.Record{
+			RequestedAt: startedAt.Add(5 * time.Millisecond), APIKey: fmt.Sprintf("key-%d", index),
+			Model: "deepseek-v4-flash", TotalTokens: int64(100 + index),
+		}); err != nil {
+			t.Fatal(err)
+		}
+		if err := store.UpsertRequest(context.Background(), collection.RequestRecord{
+			RequestID: fmt.Sprintf("request-%d", index), RequestedAt: startedAt,
+			CompletedAt: startedAt.Add(time.Second), Model: "deepseek-v4-flash", Stream: true,
+			Outcome: "succeeded", Endpoint: "/v1/responses",
+		}); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	records, err := store.ListRecent(context.Background(), 10)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(records) != 4 {
+		t.Fatalf("ambiguous records = %d, want four unmerged facts: %+v", len(records), records)
+	}
+}
+
+func TestSQLiteUsageKeepsLifecycleWithoutUsageVisible(t *testing.T) {
+	store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer store.Close()
+
+	startedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC)
+	if err := store.UpsertRequest(context.Background(), collection.RequestRecord{
+		RequestID: "canceled", RequestedAt: startedAt, CompletedAt: startedAt.Add(time.Second),
+		Model: "deepseek-v4-flash", Stream: true, Outcome: "canceled", StatusCode: 499,
+		Endpoint: "/v1/responses",
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	records, err := store.ListRecent(context.Background(), 10)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(records) != 1 || records[0].RequestID != "canceled" || records[0].Outcome != "canceled" {
+		t.Fatalf("records = %+v", records)
+	}
+}