diff --git a/internal/repository/sqlite_usage.go b/internal/repository/sqlite_usage.go index 3f0d03c..3216dff 100644 --- a/internal/repository/sqlite_usage.go +++ b/internal/repository/sqlite_usage.go @@ -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() } diff --git a/internal/repository/sqlite_usage_test.go b/internal/repository/sqlite_usage_test.go index 19b7e21..698bd6d 100644 --- a/internal/repository/sqlite_usage_test.go +++ b/internal/repository/sqlite_usage_test.go @@ -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) + } +}