From 0e3173189f79c441396153346d16b0e239240301 Mon Sep 17 00:00:00 2001
From: dark <nndaaa@gmail.com>
Date: Sat, 14 Feb 2026 22:11:20 +0800
Subject: [PATCH] feat: add AI billing module (quota + usage)

- QuotaService: CheckAndFreeze, Settle, Unfreeze
- UsageService: Record, UpdateConversationStats
- Skip billing for user-provided API keys
---
 backend/internal/ai/billing/quota.go | 50 ++++++++++++++++++++++++++++
 backend/internal/ai/billing/usage.go | 32 ++++++++++++++++++
 2 files changed, 82 insertions(+)
 create mode 100644 backend/internal/ai/billing/quota.go
 create mode 100644 backend/internal/ai/billing/usage.go

diff --git a/backend/internal/ai/billing/quota.go b/backend/internal/ai/billing/quota.go
new file mode 100644
index 0000000..21c14a1
--- /dev/null
+++ b/backend/internal/ai/billing/quota.go
@@ -0,0 +1,50 @@
+package billing
+
+import (
+	"context"
+
+	"github.com/youruser/base/model"
+
+	"gorm.io/gorm"
+)
+
+// QuotaService handles quota freeze/settle/unfreeze operations
+type QuotaService struct{}
+
+func NewQuotaService() *QuotaService {
+	return &QuotaService{}
+}
+
+// CheckAndFreeze checks if user has sufficient balance and freezes the estimated cost.
+// Returns nil if successful, error if insufficient balance.
+// If apiKeyId > 0, it's a user-provided key — skip billing.
+func (s *QuotaService) CheckAndFreeze(ctx context.Context, db *gorm.DB, userId int64, estimatedCost float64, apiKeyId int64) error {
+	// User-provided keys skip billing
+	if apiKeyId > 0 {
+		return nil
+	}
+
+	// Ensure quota record exists
+	_, err := model.AIUserQuotaEnsure(ctx, db, userId)
+	if err != nil {
+		return err
+	}
+
+	return model.AIUserQuotaFreeze(ctx, db, userId, estimatedCost)
+}
+
+// Settle finalizes billing: releases frozen amount, deducts actual cost, refunds difference.
+func (s *QuotaService) Settle(ctx context.Context, db *gorm.DB, userId int64, frozenAmount, actualCost float64, apiKeyId int64) error {
+	if apiKeyId > 0 {
+		return nil
+	}
+	return model.AIUserQuotaSettle(ctx, db, userId, frozenAmount, actualCost)
+}
+
+// Unfreeze releases frozen amount back to balance (used on error).
+func (s *QuotaService) Unfreeze(ctx context.Context, db *gorm.DB, userId int64, amount float64, apiKeyId int64) error {
+	if apiKeyId > 0 {
+		return nil
+	}
+	return model.AIUserQuotaUnfreeze(ctx, db, userId, amount)
+}
diff --git a/backend/internal/ai/billing/usage.go b/backend/internal/ai/billing/usage.go
new file mode 100644
index 0000000..5b5b9d3
--- /dev/null
+++ b/backend/internal/ai/billing/usage.go
@@ -0,0 +1,32 @@
+package billing
+
+import (
+	"context"
+
+	"github.com/youruser/base/model"
+
+	"gorm.io/gorm"
+)
+
+// UsageService handles usage recording
+type UsageService struct{}
+
+func NewUsageService() *UsageService {
+	return &UsageService{}
+}
+
+// Record inserts a usage record
+func (s *UsageService) Record(ctx context.Context, db *gorm.DB, record *model.AIUsageRecord) error {
+	_, err := model.AIUsageRecordInsert(ctx, db, record)
+	return err
+}
+
+// UpdateConversationStats updates conversation token count and cost
+func (s *UsageService) UpdateConversationStats(ctx context.Context, db *gorm.DB, conversationId int64, tokens int64, cost float64) error {
+	return db.WithContext(ctx).Model(&model.AIConversation{}).
+		Where("id = ?", conversationId).
+		Updates(map[string]interface{}{
+			"total_tokens": gorm.Expr("total_tokens + ?", tokens),
+			"total_cost":   gorm.Expr("total_cost + ?", cost),
+		}).Error
+}