package app import ( "context" "encoding/json" "errors" "strings" "time" "unicode/utf8" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache" "github.com/geo-platform/tenant-api/internal/shared/response" "github.com/geo-platform/tenant-api/internal/tenant/repository" ) const ( AIUsageTypeArticleSelectionOptimize = "article_selection_optimize" AIUsageTypeTemplateAnalyze = "template_analyze" AIUsageTypeTemplateTitleGenerate = "template_title_generate" AIUsageTypeTemplateOutlineGenerate = "template_outline_generate" AIUsageTypeKolPromptGenerate = "kol_prompt_generate" AIUsageTypeKolPromptOptimize = "kol_prompt_optimize" aiPointsQuotaType = "ai_points" ) type AIPointReserveInput struct { TenantID int64 OperatorID int64 UsageType string ResourceType *string ResourceID *int64 ResourceUID *string MeteredText string FixedPoints int Metadata map[string]any } type AIPointReservation struct { ReservationID int64 UsageLogID int64 Points int RequestChars int BaseChars int } func ReserveAIPoints( ctx context.Context, pool *pgxpool.Pool, cache sharedcache.Cache, input AIPointReserveInput, ) (*AIPointReservation, error) { if pool == nil { return nil, response.ErrServiceUnavailable(50344, "ai_points_unavailable", "ai points store is unavailable") } input.UsageType = strings.TrimSpace(input.UsageType) if input.UsageType == "" { return nil, response.ErrBadRequest(40075, "ai_points_usage_type_required", "usage type is required") } tx, err := pool.Begin(ctx) if err != nil { return nil, response.ErrInternal(50120, "ai_points_tx_failed", "failed to begin ai points transaction") } defer rollbackGenerationTx(tx) if err := lockTenantAIPoints(ctx, tx, input.TenantID); err != nil { return nil, response.ErrInternal(50121, "ai_points_lock_failed", "failed to lock ai points quota") } now := time.Now().UTC() quotaRepo := repository.NewQuotaRepository(tx) status, err := quotaRepo.GetAIQuotaStatus(ctx, input.TenantID, now) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return nil, response.ErrForbidden(40381, "ai_points_insufficient", "AI points are not available for this plan") } return nil, response.ErrInternal(50122, "ai_points_query_failed", "failed to query ai points") } requestChars := countAIPointChars(input.MeteredText) points := calculateAIPointCost(requestChars, status.BaseChars) if input.FixedPoints > 0 { points = input.FixedPoints } if status.Total <= 0 || status.Balance < points { return nil, response.ErrForbidden(40381, "ai_points_insufficient", "AI points are insufficient, please upgrade your plan") } resourceType := "ai_usage" if input.ResourceType != nil && strings.TrimSpace(*input.ResourceType) != "" { resourceType = strings.TrimSpace(*input.ResourceType) } reservationID, err := quotaRepo.CreateQuotaReservation(ctx, repository.QuotaReservationInput{ TenantID: input.TenantID, OperatorID: input.OperatorID, QuotaType: aiPointsQuotaType, ResourceType: resourceType, ResourceID: input.ResourceID, ReservedAmount: points, ExpireAt: now.Add(time.Hour), }) if err != nil { return nil, response.ErrInternal(50123, "ai_points_reserve_failed", "failed to reserve ai points") } metadata := input.Metadata if metadata == nil { metadata = map[string]any{} } metadataJSON, err := json.Marshal(metadata) if err != nil { return nil, response.ErrInternal(50124, "ai_points_metadata_failed", "failed to encode ai points metadata") } usageID, err := repository.NewAIPointUsageRepository(tx).Create(ctx, repository.CreateAIPointUsageLogInput{ TenantID: input.TenantID, OperatorID: input.OperatorID, QuotaReservationID: reservationID, UsageType: input.UsageType, ResourceType: &resourceType, ResourceID: input.ResourceID, ResourceUID: input.ResourceUID, RequestChars: requestChars, BaseChars: status.BaseChars, Points: points, MetadataJSON: metadataJSON, }) if err != nil { return nil, response.ErrInternal(50125, "ai_points_audit_failed", "failed to create ai points audit log") } reason := "ai_points_reserve" referenceType := "ai_point_usage" balanceAfter := status.Balance - points if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{ TenantID: input.TenantID, OperatorID: input.OperatorID, QuotaType: aiPointsQuotaType, Delta: -points, BalanceAfter: balanceAfter, Reason: &reason, ReferenceType: &referenceType, ReferenceID: &usageID, }); err != nil { return nil, response.ErrInternal(50126, "ai_points_ledger_failed", "failed to write ai points ledger") } if err := tx.Commit(ctx); err != nil { return nil, response.ErrInternal(50127, "ai_points_commit_failed", "failed to commit ai points reservation") } deleteCacheKey(ctx, cache, workspaceQuotaSummaryCacheKey(input.TenantID)) return &AIPointReservation{ ReservationID: reservationID, UsageLogID: usageID, Points: points, RequestChars: requestChars, BaseChars: status.BaseChars, }, nil } func CompleteAIPoints( ctx context.Context, pool *pgxpool.Pool, tenantID int64, reservation AIPointReservation, model string, ) error { if pool == nil || reservation.ReservationID <= 0 || reservation.UsageLogID <= 0 { return nil } tx, err := pool.Begin(ctx) if err != nil { return err } defer rollbackGenerationTx(tx) pending, err := lockPendingAIPointUsage(ctx, tx, tenantID, reservation.UsageLogID) if err != nil || !pending { return err } if err := confirmAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil { return err } var modelPtr *string if strings.TrimSpace(model) != "" { normalized := strings.TrimSpace(model) modelPtr = &normalized } if err := repository.NewAIPointUsageRepository(tx).MarkCompleted(ctx, tenantID, reservation.UsageLogID, modelPtr); err != nil { return err } return tx.Commit(ctx) } func RefundAIPoints( ctx context.Context, pool *pgxpool.Pool, cache sharedcache.Cache, tenantID int64, operatorID int64, reservation AIPointReservation, errorMessage string, ) error { if pool == nil || reservation.ReservationID <= 0 || reservation.UsageLogID <= 0 || reservation.Points <= 0 { return nil } tx, err := pool.Begin(ctx) if err != nil { return err } defer rollbackGenerationTx(tx) if err := lockTenantAIPoints(ctx, tx, tenantID); err != nil { return err } pending, err := lockPendingAIPointUsage(ctx, tx, tenantID, reservation.UsageLogID) if err != nil || !pending { return err } status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC()) if err != nil { return err } quotaRepo := repository.NewQuotaRepository(tx) reason := "ai_points_refund" referenceType := "ai_point_usage" usageID := reservation.UsageLogID if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{ TenantID: tenantID, OperatorID: operatorID, QuotaType: aiPointsQuotaType, Delta: reservation.Points, BalanceAfter: status.Balance + reservation.Points, Reason: &reason, ReferenceType: &referenceType, ReferenceID: &usageID, }); err != nil { return err } if err := refundAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil { return err } if err := repository.NewAIPointUsageRepository(tx).MarkRefunded(ctx, tenantID, reservation.UsageLogID, strings.TrimSpace(errorMessage)); err != nil { return err } if err := tx.Commit(ctx); err != nil { return err } deleteCacheKey(ctx, cache, workspaceQuotaSummaryCacheKey(tenantID)) return nil } func CompletePendingAIPointsByResource( ctx context.Context, pool *pgxpool.Pool, tenantID int64, resourceType string, resourceID *int64, resourceUID *string, model string, ) error { reservation, err := findPendingAIPointReservation(ctx, pool, tenantID, resourceType, resourceID, resourceUID) if err != nil || reservation == nil { return err } return CompleteAIPoints(ctx, pool, tenantID, *reservation, model) } func RefundPendingAIPointsByResource( ctx context.Context, pool *pgxpool.Pool, cache sharedcache.Cache, tenantID int64, resourceType string, resourceID *int64, resourceUID *string, errorMessage string, ) error { reservation, operatorID, err := findPendingAIPointReservationWithOperator(ctx, pool, tenantID, resourceType, resourceID, resourceUID) if err != nil || reservation == nil { return err } return RefundAIPoints(ctx, pool, cache, tenantID, operatorID, *reservation, errorMessage) } func findPendingAIPointReservation( ctx context.Context, pool *pgxpool.Pool, tenantID int64, resourceType string, resourceID *int64, resourceUID *string, ) (*AIPointReservation, error) { reservation, _, err := findPendingAIPointReservationWithOperator(ctx, pool, tenantID, resourceType, resourceID, resourceUID) return reservation, err } func findPendingAIPointReservationWithOperator( ctx context.Context, pool *pgxpool.Pool, tenantID int64, resourceType string, resourceID *int64, resourceUID *string, ) (*AIPointReservation, int64, error) { if pool == nil { return nil, 0, nil } log, err := repository.NewAIPointUsageRepository(pool).FindPendingByResource(ctx, tenantID, resourceType, resourceID, resourceUID) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return nil, 0, nil } return nil, 0, err } if log == nil || log.QuotaReservationID == nil { return nil, 0, nil } operatorID := int64(0) if log.OperatorID != nil { operatorID = *log.OperatorID } return &AIPointReservation{ ReservationID: *log.QuotaReservationID, UsageLogID: log.ID, Points: log.Points, RequestChars: log.RequestChars, BaseChars: log.BaseChars, }, operatorID, nil } func countAIPointChars(text string) int { return utf8.RuneCountInString(strings.TrimSpace(text)) } func calculateAIPointCost(chars, baseChars int) int { if baseChars <= 0 { baseChars = 1000 } if chars <= 0 { return 1 } return chars/baseChars + 1 } func lockTenantAIPoints(ctx context.Context, tx pgx.Tx, tenantID int64) error { _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(20260428, hashint8($1)::int)`, tenantID) return err } func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogID int64) (bool, error) { var status string err := tx.QueryRow(ctx, ` SELECT status FROM ai_point_usage_logs WHERE id = $1 AND tenant_id = $2 FOR UPDATE `, usageLogID, tenantID).Scan(&status) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return false, nil } return false, err } return status == "pending", nil } func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { tag, err := tx.Exec(ctx, ` UPDATE quota_reservations SET status = 'confirmed', consumed_amount = reserved_amount, updated_at = NOW() WHERE id = $1 AND tenant_id = $2 AND quota_type = $3 AND status = 'pending' `, reservationID, tenantID, aiPointsQuotaType) if err != nil { return err } if tag.RowsAffected() == 0 { return errors.New("ai points reservation is not pending") } return nil } func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { tag, err := tx.Exec(ctx, ` UPDATE quota_reservations SET status = 'refunded', refunded_amount = reserved_amount, updated_at = NOW() WHERE id = $1 AND tenant_id = $2 AND quota_type = $3 AND status = 'pending' `, reservationID, tenantID, aiPointsQuotaType) if err != nil { return err } if tag.RowsAffected() == 0 { return errors.New("ai points reservation is not pending") } return nil }