feat(tenant): bill AI points by output characters
Deployment Config CI / Deployment Config (push) Successful in 39s
Backend CI / Backend (push) Successful in 18m14s

Previously AI point cost was computed from request (input) characters at
reservation time. Switch to reserving a single point up front and
settling the final cost from the generated output length on completion.

- Reserve 1 point (or FixedPoints) instead of pricing on request chars
- On completion, recompute points from output chars vs base chars,
  update the reservation amounts, and post a ledger delta for the
  top-up/refund; invalidate the workspace quota summary cache
- MarkCompleted now persists final request_chars/base_chars/points
- Use ceil division in calculateAIPointCost so an exact base boundary
  charges one point instead of rolling over
- Thread output text through all CompleteAIPoints callers (article
  selection, KOL assist, question expansion, template assist, compliance
  judge) and return the settled reservation
- Add unit tests plus an integration test gated on TEST_DATABASE_URL
- Update the user manual: billing is by output characters

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-07-09 17:04:37 +08:00
parent c8dceb43d7
commit f69edc2218
12 changed files with 612 additions and 51 deletions
+117 -17
View File
@@ -44,11 +44,20 @@ type AIPointReserveInput struct {
type AIPointReservation struct {
ReservationID int64
UsageLogID int64
OperatorID int64
Points int
RequestChars int
BaseChars int
}
type aiPointOutputSettlement struct {
OutputChars int
BaseChars int
FinalPoints int
ReservedPoints int
LedgerDelta int
}
func ReserveAIPoints(
ctx context.Context,
pool *pgxpool.Pool,
@@ -85,8 +94,8 @@ func ReserveAIPoints(
}
requestChars := countAIPointChars(input.MeteredText)
points := calculateAIPointCost(requestChars, status.BaseChars)
if input.FixedPoints > 0 {
points := 1
if input.FixedPoints > points {
points = input.FixedPoints
}
if status.Total <= 0 || status.Balance < points {
@@ -160,6 +169,7 @@ func ReserveAIPoints(
return &AIPointReservation{
ReservationID: reservationID,
UsageLogID: usageID,
OperatorID: input.OperatorID,
Points: points,
RequestChars: requestChars,
BaseChars: status.BaseChars,
@@ -169,36 +179,79 @@ func ReserveAIPoints(
func CompleteAIPoints(
ctx context.Context,
pool *pgxpool.Pool,
cache sharedcache.Cache,
tenantID int64,
reservation AIPointReservation,
model string,
) error {
meteredOutputText string,
) (*AIPointReservation, error) {
completed := reservation
if pool == nil || reservation.ReservationID <= 0 || reservation.UsageLogID <= 0 {
return nil
return &completed, nil
}
tx, err := pool.Begin(ctx)
if err != nil {
return err
return nil, err
}
defer rollbackGenerationTx(tx)
pending, err := lockPendingAIPointUsage(ctx, tx, tenantID, reservation.UsageLogID)
if err != nil || !pending {
return err
if err := lockTenantAIPoints(ctx, tx, tenantID); err != nil {
return nil, err
}
if err := confirmAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID); err != nil {
return err
pending, err := lockPendingAIPointUsage(ctx, tx, tenantID, reservation.UsageLogID)
if err != nil {
return nil, err
}
if !pending {
return &completed, nil
}
settlement := calculateAIPointOutputSettlement(reservation, meteredOutputText)
completed.RequestChars = settlement.OutputChars
completed.BaseChars = settlement.BaseChars
completed.Points = settlement.FinalPoints
quotaRepo := repository.NewQuotaRepository(tx)
adjusted, err := completeAIPointReservationWithFinalAmountIfPending(ctx, tx, tenantID, reservation.ReservationID, settlement.FinalPoints)
if err != nil {
return nil, err
}
if adjusted && settlement.LedgerDelta != 0 {
currentBalance, err := quotaRepo.GetCurrentBalance(ctx, tenantID, aiPointsQuotaType)
if err != nil {
return nil, err
}
reason := "ai_points_output_settle"
referenceType := "ai_point_usage"
usageID := reservation.UsageLogID
if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
TenantID: tenantID,
OperatorID: reservation.OperatorID,
QuotaType: aiPointsQuotaType,
Delta: settlement.LedgerDelta,
BalanceAfter: currentBalance + settlement.LedgerDelta,
Reason: &reason,
ReferenceType: &referenceType,
ReferenceID: &usageID,
}); err != nil {
return nil, 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
if err := repository.NewAIPointUsageRepository(tx).MarkCompleted(ctx, tenantID, reservation.UsageLogID, settlement.OutputChars, settlement.BaseChars, settlement.FinalPoints, modelPtr); err != nil {
return nil, err
}
return tx.Commit(ctx)
if err := tx.Commit(ctx); err != nil {
return nil, err
}
deleteCacheKey(ctx, cache, workspaceQuotaSummaryCacheKey(tenantID))
return &completed, nil
}
func RefundAIPoints(
@@ -267,17 +320,19 @@ func RefundAIPoints(
func CompletePendingAIPointsByResource(
ctx context.Context,
pool *pgxpool.Pool,
cache sharedcache.Cache,
tenantID int64,
resourceType string,
resourceID *int64,
resourceUID *string,
model string,
) error {
meteredOutputText string,
) (*AIPointReservation, error) {
reservation, err := findPendingAIPointReservation(ctx, pool, tenantID, resourceType, resourceID, resourceUID)
if err != nil || reservation == nil {
return err
return reservation, err
}
return CompleteAIPoints(ctx, pool, tenantID, *reservation, model)
return CompleteAIPoints(ctx, pool, cache, tenantID, *reservation, model, meteredOutputText)
}
func RefundPendingAIPointsByResource(
@@ -337,6 +392,7 @@ func findPendingAIPointReservationWithOperator(
return &AIPointReservation{
ReservationID: *log.QuotaReservationID,
UsageLogID: log.ID,
OperatorID: operatorID,
Points: log.Points,
RequestChars: log.RequestChars,
BaseChars: log.BaseChars,
@@ -354,7 +410,27 @@ func calculateAIPointCost(chars, baseChars int) int {
if chars <= 0 {
return 1
}
return chars/baseChars + 1
return (chars + baseChars - 1) / baseChars
}
func calculateAIPointOutputSettlement(reservation AIPointReservation, meteredOutputText string) aiPointOutputSettlement {
baseChars := reservation.BaseChars
if baseChars <= 0 {
baseChars = 1000
}
outputChars := countAIPointChars(meteredOutputText)
finalPoints := calculateAIPointCost(outputChars, baseChars)
reservedPoints := reservation.Points
if reservedPoints <= 0 {
reservedPoints = 1
}
return aiPointOutputSettlement{
OutputChars: outputChars,
BaseChars: baseChars,
FinalPoints: finalPoints,
ReservedPoints: reservedPoints,
LedgerDelta: reservedPoints - finalPoints,
}
}
func lockTenantAIPoints(ctx context.Context, tx pgx.Tx, tenantID int64) error {
@@ -379,6 +455,30 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI
return status == "pending", nil
}
func completeAIPointReservationWithFinalAmountIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64, finalPoints int) (bool, error) {
if finalPoints <= 0 {
finalPoints = 1
}
tag, err := tx.Exec(ctx, `
UPDATE quota_reservations
SET status = 'confirmed',
reserved_amount = $4,
consumed_amount = $4,
updated_at = NOW()
WHERE id = $1
AND tenant_id = $2
AND quota_type = $3
AND status = 'pending'
`, reservationID, tenantID, aiPointsQuotaType, finalPoints)
if err != nil {
return false, err
}
if tag.RowsAffected() > 0 {
return true, nil
}
return false, ignoreAIPointReservationReset(ctx, tx, false, tenantID, reservationID)
}
func confirmAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error {
tag, err := tx.Exec(ctx, `
UPDATE quota_reservations
@@ -0,0 +1,200 @@
//go:build integration
package app
import (
"context"
"fmt"
"os"
"strings"
"testing"
"time"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/geo-platform/tenant-api/internal/tenant/repository"
)
func TestAIPointsOutputSettlementRealDatabase(t *testing.T) {
dsn := os.Getenv("TEST_DATABASE_URL")
if dsn == "" {
t.Skip("set TEST_DATABASE_URL to run the real database AI point settlement test")
}
ctx := context.Background()
pool, err := pgxpool.New(ctx, dsn)
if err != nil {
t.Fatalf("connect database: %v", err)
}
t.Cleanup(pool.Close)
tenantID, userID := seedAIPointSettlementTenant(t, ctx, pool, 10, 1000)
reservation, err := ReserveAIPoints(ctx, pool, nil, AIPointReserveInput{
TenantID: tenantID,
OperatorID: userID,
UsageType: "test_output_billing",
MeteredText: "cheap input",
})
if err != nil {
t.Fatalf("reserve ai points: %v", err)
}
if reservation.Points != 1 {
t.Fatalf("reserved points = %d, want minimum 1", reservation.Points)
}
output := strings.Repeat("贵", 2500)
completed, err := CompleteAIPoints(ctx, pool, nil, tenantID, *reservation, "test-model", output)
if err != nil {
t.Fatalf("complete ai points: %v", err)
}
if completed.Points != 3 {
t.Fatalf("completed points = %d, want 3", completed.Points)
}
if completed.RequestChars != 2500 {
t.Fatalf("completed request/output chars = %d, want 2500", completed.RequestChars)
}
var usageStatus, model string
var usageChars, usagePoints int
if err := pool.QueryRow(ctx, `
SELECT status, request_chars, points, model
FROM ai_point_usage_logs
WHERE id = $1 AND tenant_id = $2
`, reservation.UsageLogID, tenantID).Scan(&usageStatus, &usageChars, &usagePoints, &model); err != nil {
t.Fatalf("query usage log: %v", err)
}
if usageStatus != "completed" || usageChars != 2500 || usagePoints != 3 || model != "test-model" {
t.Fatalf("usage log = status %q chars %d points %d model %q; want completed 2500 3 test-model", usageStatus, usageChars, usagePoints, model)
}
var reservationStatus string
var reservedAmount, consumedAmount int
if err := pool.QueryRow(ctx, `
SELECT status, reserved_amount, consumed_amount
FROM quota_reservations
WHERE id = $1 AND tenant_id = $2
`, reservation.ReservationID, tenantID).Scan(&reservationStatus, &reservedAmount, &consumedAmount); err != nil {
t.Fatalf("query quota reservation: %v", err)
}
if reservationStatus != "confirmed" || reservedAmount != 3 || consumedAmount != 3 {
t.Fatalf("reservation = status %q reserved %d consumed %d; want confirmed 3 3", reservationStatus, reservedAmount, consumedAmount)
}
rows, err := pool.Query(ctx, `
SELECT delta, balance_after, reason
FROM tenant_quota_ledgers
WHERE tenant_id = $1
AND quota_type = 'ai_points'
AND reference_id = $2
ORDER BY id
`, tenantID, reservation.UsageLogID)
if err != nil {
t.Fatalf("query quota ledger: %v", err)
}
defer rows.Close()
var deltas []int
var balances []int
var reasons []string
for rows.Next() {
var delta, balance int
var reason string
if err := rows.Scan(&delta, &balance, &reason); err != nil {
t.Fatalf("scan quota ledger: %v", err)
}
deltas = append(deltas, delta)
balances = append(balances, balance)
reasons = append(reasons, reason)
}
if err := rows.Err(); err != nil {
t.Fatalf("iterate quota ledger: %v", err)
}
if fmt.Sprint(deltas) != "[-1 -2]" || fmt.Sprint(balances) != "[9 7]" || fmt.Sprint(reasons) != "[ai_points_reserve ai_points_output_settle]" {
t.Fatalf("ledger deltas=%v balances=%v reasons=%v; want [-1 -2], [9 7], reserve+settle", deltas, balances, reasons)
}
status, err := repository.NewQuotaRepository(pool).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC())
if err != nil {
t.Fatalf("query quota status: %v", err)
}
if status.Used != 3 || status.Balance != 7 {
t.Fatalf("quota status used=%d balance=%d; want 3/7", status.Used, status.Balance)
}
}
func seedAIPointSettlementTenant(t *testing.T, ctx context.Context, pool *pgxpool.Pool, aiPointsMonthly int, baseChars int) (int64, int64) {
t.Helper()
suffix := fmt.Sprintf("%d", time.Now().UnixNano())
var tenantID int64
if err := pool.QueryRow(ctx, `
INSERT INTO tenants (name, status)
VALUES ($1, 'active')
RETURNING id
`, "ai-point-output-test-"+suffix).Scan(&tenantID); err != nil {
t.Fatalf("insert tenant: %v", err)
}
var userID int64
if err := pool.QueryRow(ctx, `
INSERT INTO users (phone, password_hash, name, status)
VALUES ($1, 'test', 'AI Point Output Test', 'active')
RETURNING id
`, "199"+suffix[len(suffix)-8:]).Scan(&userID); err != nil {
t.Fatalf("insert user: %v", err)
}
if _, err := pool.Exec(ctx, `
INSERT INTO tenant_memberships (tenant_id, user_id, tenant_role)
VALUES ($1, $2, 'owner')
`, tenantID, userID); err != nil {
t.Fatalf("insert tenant membership: %v", err)
}
var planID int64
if err := pool.QueryRow(ctx, `
INSERT INTO plans (plan_code, name, quota_policy_json, status)
VALUES ($1, 'AI Point Output Plan', jsonb_build_object(
'article_generation', 0,
'article_quota_cycle', 'monthly',
'ai_points_monthly', $2::int,
'ai_point_base_chars', $3::int,
'image_storage_bytes', 0,
'brand_limit', 1
), 'active')
RETURNING id
`, "ai-point-output-"+suffix, aiPointsMonthly, baseChars).Scan(&planID); err != nil {
t.Fatalf("insert plan: %v", err)
}
if _, err := pool.Exec(ctx, `
INSERT INTO tenant_plan_subscriptions (tenant_id, plan_id, start_at, end_at, status)
VALUES ($1, $2, NOW() - INTERVAL '1 hour', NOW() + INTERVAL '24 hours', 'active')
`, tenantID, planID); err != nil {
t.Fatalf("insert tenant plan subscription: %v", err)
}
t.Cleanup(func() {
cleanupCtx := context.Background()
mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM ai_point_usage_logs WHERE tenant_id = $1`, tenantID)
mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM tenant_quota_ledgers WHERE tenant_id = $1`, tenantID)
mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM quota_reservations WHERE tenant_id = $1`, tenantID)
mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM tenant_plan_subscriptions WHERE tenant_id = $1`, tenantID)
mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM workspace_memberships WHERE tenant_id = $1`, tenantID)
mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM workspaces WHERE tenant_id = $1`, tenantID)
mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM tenant_memberships WHERE tenant_id = $1`, tenantID)
mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM users WHERE id = $1`, userID)
mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM tenants WHERE id = $1`, tenantID)
mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM plans WHERE id = $1`, planID)
})
return tenantID, userID
}
func mustCleanupAIPointSettlementTestData(t *testing.T, ctx context.Context, pool *pgxpool.Pool, sql string, args ...any) {
t.Helper()
if _, err := pool.Exec(ctx, sql, args...); err != nil {
t.Errorf("cleanup AI point settlement test data: %v; sql=%s", err, sql)
}
}
@@ -1,6 +1,14 @@
package app
import "testing"
import (
"context"
"errors"
"strings"
"testing"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
)
func TestCalculateAIPointCost(t *testing.T) {
t.Parallel()
@@ -13,12 +21,12 @@ func TestCalculateAIPointCost(t *testing.T) {
}{
{name: "empty still charges base usage", chars: 0, baseChars: 1000, want: 1},
{name: "single char", chars: 1, baseChars: 1000, want: 1},
{name: "first base boundary", chars: 1000, baseChars: 1000, want: 2},
{name: "first base boundary", chars: 1000, baseChars: 1000, want: 1},
{name: "over first base", chars: 1001, baseChars: 1000, want: 2},
{name: "second base boundary", chars: 2000, baseChars: 1000, want: 3},
{name: "second base boundary", chars: 2000, baseChars: 1000, want: 2},
{name: "over second base", chars: 2001, baseChars: 1000, want: 3},
{name: "default base chars", chars: 1000, baseChars: 0, want: 2},
{name: "custom base chars", chars: 1500, baseChars: 500, want: 4},
{name: "default base chars", chars: 1000, baseChars: 0, want: 1},
{name: "custom base chars", chars: 1500, baseChars: 500, want: 3},
}
for _, tt := range tests {
@@ -31,3 +39,150 @@ func TestCalculateAIPointCost(t *testing.T) {
})
}
}
func TestCalculateAIPointOutputSettlementUsesOutputText(t *testing.T) {
t.Parallel()
reservation := AIPointReservation{
Points: 1,
RequestChars: 8,
BaseChars: 1000,
}
got := calculateAIPointOutputSettlement(reservation, strings.Repeat("贵", 2500))
if got.OutputChars != 2500 {
t.Fatalf("OutputChars = %d, want 2500", got.OutputChars)
}
if got.FinalPoints != 3 {
t.Fatalf("FinalPoints = %d, want 3", got.FinalPoints)
}
if got.LedgerDelta != -2 {
t.Fatalf("LedgerDelta = %d, want -2 for output-side top-up", got.LedgerDelta)
}
}
func TestCalculateAIPointOutputSettlementBoundaries(t *testing.T) {
t.Parallel()
tests := []struct {
name string
outputChars int
wantPoints int
}{
{name: "minimum empty output still one point", outputChars: 0, wantPoints: 1},
{name: "below base", outputChars: 999, wantPoints: 1},
{name: "exactly base", outputChars: 1000, wantPoints: 1},
{name: "one over base", outputChars: 1001, wantPoints: 2},
{name: "exactly second base", outputChars: 2000, wantPoints: 2},
{name: "one over second base", outputChars: 2001, wantPoints: 3},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
got := calculateAIPointOutputSettlement(
AIPointReservation{Points: 1, RequestChars: 50000, BaseChars: 1000},
strings.Repeat("字", tt.outputChars),
)
if got.FinalPoints != tt.wantPoints {
t.Fatalf("FinalPoints = %d, want %d", got.FinalPoints, tt.wantPoints)
}
})
}
}
func TestCalculateAIPointOutputSettlementRefundsOverReservation(t *testing.T) {
t.Parallel()
got := calculateAIPointOutputSettlement(
AIPointReservation{Points: 5, RequestChars: 5000, BaseChars: 1000},
strings.Repeat("字", 700),
)
if got.FinalPoints != 1 {
t.Fatalf("FinalPoints = %d, want 1", got.FinalPoints)
}
if got.LedgerDelta != 4 {
t.Fatalf("LedgerDelta = %d, want +4 for output-side refund", got.LedgerDelta)
}
}
func TestCalculateAIPointOutputSettlementUsesDefaultBaseChars(t *testing.T) {
t.Parallel()
got := calculateAIPointOutputSettlement(
AIPointReservation{Points: 1, BaseChars: 0},
strings.Repeat("字", 1001),
)
if got.BaseChars != 1000 {
t.Fatalf("BaseChars = %d, want default 1000", got.BaseChars)
}
if got.FinalPoints != 2 {
t.Fatalf("FinalPoints = %d, want 2", got.FinalPoints)
}
if got.LedgerDelta != -1 {
t.Fatalf("LedgerDelta = %d, want -1", got.LedgerDelta)
}
}
func TestCompleteAIPointReservationWithFinalAmountUpdatesReservationAmounts(t *testing.T) {
tx := &captureAIPointTx{}
adjusted, err := completeAIPointReservationWithFinalAmountIfPending(context.Background(), tx, 10, 20, 3)
if err != nil {
t.Fatalf("complete reservation: %v", err)
}
if !adjusted {
t.Fatalf("adjusted = false, want true")
}
if !strings.Contains(tx.sql, "reserved_amount = $4") {
t.Fatalf("reservation SQL must set final reserved_amount; query:\n%s", tx.sql)
}
if !strings.Contains(tx.sql, "consumed_amount = $4") {
t.Fatalf("reservation SQL must set final consumed_amount; query:\n%s", tx.sql)
}
if len(tx.args) != 4 {
t.Fatalf("args length = %d, want 4", len(tx.args))
}
if tx.args[0] != int64(20) || tx.args[1] != int64(10) || tx.args[2] != aiPointsQuotaType || tx.args[3] != 3 {
t.Fatalf("args = %#v, want reservation, tenant, quota type, final points", tx.args)
}
}
type captureAIPointTx struct {
sql string
args []any
}
func (tx *captureAIPointTx) Begin(context.Context) (pgx.Tx, error) { return nil, errors.New("unused") }
func (tx *captureAIPointTx) Commit(context.Context) error { return nil }
func (tx *captureAIPointTx) Rollback(context.Context) error { return nil }
func (tx *captureAIPointTx) CopyFrom(context.Context, pgx.Identifier, []string, pgx.CopyFromSource) (int64, error) {
return 0, errors.New("unused")
}
func (tx *captureAIPointTx) SendBatch(context.Context, *pgx.Batch) pgx.BatchResults {
return nil
}
func (tx *captureAIPointTx) LargeObjects() pgx.LargeObjects { return pgx.LargeObjects{} }
func (tx *captureAIPointTx) Prepare(context.Context, string, string) (*pgconn.StatementDescription, error) {
return nil, errors.New("unused")
}
func (tx *captureAIPointTx) Exec(_ context.Context, sql string, args ...any) (pgconn.CommandTag, error) {
tx.sql = sql
tx.args = args
return pgconn.NewCommandTag("UPDATE 1"), nil
}
func (tx *captureAIPointTx) Query(context.Context, string, ...any) (pgx.Rows, error) {
return nil, errors.New("unused")
}
func (tx *captureAIPointTx) QueryRow(context.Context, string, ...any) pgx.Row {
return failingAIPointRow{}
}
func (tx *captureAIPointTx) Conn() *pgx.Conn { return nil }
type failingAIPointRow struct{}
func (failingAIPointRow) Scan(...any) error { return errors.New("unused") }
@@ -135,7 +135,7 @@ func (s *ArticleSelectionOptimizeService) OptimizeSelection(
}
completeCtx, cancel := newGenerationCleanupContext()
if err := CompleteAIPoints(completeCtx, s.pool, actor.TenantID, *reservation, result.Model); err != nil {
if _, err := CompleteAIPoints(completeCtx, s.pool, s.cache, actor.TenantID, *reservation, result.Model, content); err != nil {
cancel()
return nil, response.ErrInternal(50019, "article_selection_optimize_failed", "failed to confirm ai points usage")
}
@@ -77,7 +77,7 @@ func (s *KolAssistService) Run(
}
completeCtx, cancel := newGenerationCleanupContext()
if err := CompleteAIPoints(completeCtx, s.pool, actor.TenantID, *reservation, generated.Model); err != nil {
if _, err := CompleteAIPoints(completeCtx, s.pool, s.cache, actor.TenantID, *reservation, generated.Model, generated.Content); err != nil {
cancel()
return nil, "", response.ErrInternal(50086, "kol_assist_result_invalid", "failed to confirm ai points usage")
}
@@ -317,7 +317,8 @@ func (s *QuestionExpansionService) GenerateByAIDistill(ctx context.Context, bran
candidates = append(candidates, buildCandidate(text, QuestionSourceAIDistill, false, existing, brandCtx))
}
if err := CompleteAIPoints(context.Background(), s.pool, actor.TenantID, *reservation, result.Model); err != nil {
completed, err := CompleteAIPoints(context.Background(), s.pool, s.cache, actor.TenantID, *reservation, result.Model, result.Content)
if err != nil {
return nil, response.ErrInternal(50127, "ai_points_commit_failed", "failed to confirm ai points")
}
@@ -329,7 +330,7 @@ func (s *QuestionExpansionService) GenerateByAIDistill(ctx context.Context, bran
return &QuestionCandidateResult{
Candidates: candidates,
AIPointsCharged: reservation.Points,
AIPointsCharged: completed.Points,
}, nil
}
@@ -413,10 +414,11 @@ func (s *QuestionExpansionService) FillQuestionCombination(ctx context.Context,
}
filled := sanitizeQuestionCombinationFillResult(payload, fallback)
if err := CompleteAIPoints(context.Background(), s.pool, actor.TenantID, *reservation, result.Model); err != nil {
completed, err := CompleteAIPoints(context.Background(), s.pool, s.cache, actor.TenantID, *reservation, result.Model, result.Content)
if err != nil {
return nil, response.ErrInternal(50127, "ai_points_commit_failed", "failed to confirm ai points")
}
filled.AIPointsCharged = reservation.Points
filled.AIPointsCharged = completed.Points
if s.cache != nil {
cached := filled
cached.AIPointsCharged = 0
@@ -430,7 +430,7 @@ func (s *TemplateService) IsAssistTaskTerminal(ctx context.Context, job AssistJo
}
switch record.Status {
case "completed":
s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, "")
s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, "", string(record.ResultJSON))
return true, nil
case "failed":
if isTemplateAssistAIPointTask(job.TaskType) {
@@ -479,7 +479,7 @@ func (s *TemplateService) executeAnalyzeTask(ctx context.Context, repo repositor
s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, err)
return
}
s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model)
s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model, result.Content)
}
func (s *TemplateService) executeTitleTask(ctx context.Context, repo repository.TemplateAssistTaskRepository, job assistJob) {
@@ -518,7 +518,7 @@ func (s *TemplateService) executeTitleTask(ctx context.Context, repo repository.
s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, err)
return
}
s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model)
s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model, result.Content)
}
func (s *TemplateService) executeOutlineTask(ctx context.Context, repo repository.TemplateAssistTaskRepository, job assistJob) {
@@ -562,7 +562,7 @@ func (s *TemplateService) executeOutlineTask(ctx context.Context, repo repositor
s.failAssist(ctx, job.TaskID, job.TaskType, job.TenantID, err)
return
}
s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model)
s.completeTemplateAssistAIPoints(ctx, job.TenantID, job.TaskID, result.Model, result.Content)
}
func (s *TemplateService) failAssist(ctx context.Context, taskID, taskType string, tenantID int64, err error) {
@@ -619,10 +619,10 @@ func (s *TemplateService) refundReservedAIPoints(ctx context.Context, tenantID,
_ = RefundAIPoints(ctx, s.pool, s.cache, tenantID, operatorID, *reservation, errorMessage)
}
func (s *TemplateService) completeTemplateAssistAIPoints(ctx context.Context, tenantID int64, taskID string, model string) {
func (s *TemplateService) completeTemplateAssistAIPoints(ctx context.Context, tenantID int64, taskID string, model string, meteredOutputText string) {
resourceType := "template_assist_task"
resourceUID := taskID
_ = CompletePendingAIPointsByResource(ctx, s.pool, tenantID, resourceType, nil, &resourceUID, model)
_, _ = CompletePendingAIPointsByResource(ctx, s.pool, s.cache, tenantID, resourceType, nil, &resourceUID, model, meteredOutputText)
}
func isTemplateAssistAIPointTask(taskType string) bool {
+70 -6
View File
@@ -1139,7 +1139,7 @@ func (s *Service) ProcessReviewJob(ctx context.Context, msg ReviewJobMessage) er
_ = s.refundComplianceLLMJudgePoints(context.Background(), job, reservation, err)
return s.failReviewJob(ctx, job.ID, cfg.ReviewMaxAttempts, err)
}
if err := s.completeComplianceLLMJudgePoints(ctx, job, reservation, result.Model); err != nil && s.logger != nil {
if err := s.completeComplianceLLMJudgePoints(ctx, job, reservation, result.Model, result.Content); err != nil && s.logger != nil {
s.logger.Warn("compliance llm billing completion failed",
zap.Error(err),
zap.Int64("tenant_id", job.TenantID),
@@ -1348,7 +1348,7 @@ func (s *Service) reserveComplianceLLMJudgePoints(ctx context.Context, job revie
return nil, err
}
requestChars := countAIPointChars(meteredText)
points := calculateAIPointCost(requestChars, status.BaseChars)
points := 1
if status.Total <= 0 || status.Balance < points {
return nil, response.ErrForbidden(40381, "ai_points_insufficient", "AI points are insufficient")
}
@@ -1416,7 +1416,7 @@ func (s *Service) reserveComplianceLLMJudgePoints(ctx context.Context, job revie
}, nil
}
func (s *Service) completeComplianceLLMJudgePoints(ctx context.Context, job reviewJob, reservation *aiPointReservation, model string) error {
func (s *Service) completeComplianceLLMJudgePoints(ctx context.Context, job reviewJob, reservation *aiPointReservation, model string, meteredOutputText string) error {
if s == nil || s.pool == nil || reservation == nil || reservation.ReservationID <= 0 || reservation.UsageLogID <= 0 {
return nil
}
@@ -1425,19 +1425,59 @@ func (s *Service) completeComplianceLLMJudgePoints(ctx context.Context, job revi
return err
}
defer rollbackTx(tx)
if err := lockTenantAIPoints(ctx, tx, job.TenantID); err != nil {
return err
}
pending, err := lockPendingAIPointUsage(ctx, tx, job.TenantID, reservation.UsageLogID)
if err != nil || !pending {
return err
}
if err := confirmAIPointReservationIfPending(ctx, tx, job.TenantID, reservation.ReservationID); err != nil {
baseChars := reservation.BaseChars
if baseChars <= 0 {
baseChars = 1000
}
outputChars := countAIPointChars(meteredOutputText)
finalPoints := calculateAIPointCost(outputChars, baseChars)
reservedPoints := reservation.Points
if reservedPoints <= 0 {
reservedPoints = 1
}
quotaRepo := tenantrepo.NewQuotaRepository(tx)
adjusted, err := completeAIPointReservationWithFinalAmountIfPending(ctx, tx, job.TenantID, reservation.ReservationID, finalPoints)
if err != nil {
return err
}
pointDelta := finalPoints - reservedPoints
if adjusted && pointDelta != 0 {
currentBalance, err := quotaRepo.GetCurrentBalance(ctx, job.TenantID, aiPointsQuotaType)
if err != nil {
return err
}
ledgerDelta := -pointDelta
reason := "ai_points_output_settle"
referenceType := "ai_point_usage"
usageID := reservation.UsageLogID
if _, err := quotaRepo.InsertQuotaLedger(ctx, tenantrepo.QuotaLedgerInput{
TenantID: job.TenantID,
OperatorID: job.CheckedBy,
QuotaType: aiPointsQuotaType,
Delta: ledgerDelta,
BalanceAfter: currentBalance + ledgerDelta,
Reason: &reason,
ReferenceType: &referenceType,
ReferenceID: &usageID,
}); err != nil {
return err
}
}
var modelPtr *string
if strings.TrimSpace(model) != "" {
normalized := strings.TrimSpace(model)
modelPtr = &normalized
}
if err := tenantrepo.NewAIPointUsageRepository(tx).MarkCompleted(ctx, job.TenantID, reservation.UsageLogID, modelPtr); err != nil {
if err := tenantrepo.NewAIPointUsageRepository(tx).MarkCompleted(ctx, job.TenantID, reservation.UsageLogID, outputChars, baseChars, finalPoints, modelPtr); err != nil {
return err
}
return tx.Commit(ctx)
@@ -1502,7 +1542,7 @@ func calculateAIPointCost(chars, baseChars int) int {
if chars <= 0 {
return 1
}
return chars/baseChars + 1
return (chars + baseChars - 1) / baseChars
}
func lockTenantAIPoints(ctx context.Context, tx pgx.Tx, tenantID int64) error {
@@ -1527,6 +1567,30 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI
return status == "pending", nil
}
func completeAIPointReservationWithFinalAmountIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64, finalPoints int) (bool, error) {
if finalPoints <= 0 {
finalPoints = 1
}
tag, err := tx.Exec(ctx, `
UPDATE quota_reservations
SET status = 'confirmed',
reserved_amount = $4,
consumed_amount = $4,
updated_at = NOW()
WHERE id = $1
AND tenant_id = $2
AND quota_type = $3
AND status = 'pending'
`, reservationID, tenantID, aiPointsQuotaType, finalPoints)
if err != nil {
return false, err
}
if tag.RowsAffected() > 0 {
return true, nil
}
return false, ignoreAIPointReservationReset(ctx, tx, false, tenantID, reservationID)
}
func confirmAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error {
tag, err := tx.Exec(ctx, `
UPDATE quota_reservations
@@ -1,6 +1,7 @@
package compliance
import (
"strings"
"testing"
sharedcompliance "github.com/geo-platform/tenant-api/internal/shared/compliance"
@@ -34,3 +35,39 @@ func TestDecidePassesWithoutHits(t *testing.T) {
t.Fatalf("decide without hits = %s, want %s", got, sharedcompliance.GateDecisionPass)
}
}
func TestComplianceAIPointCostUsesOutputBoundary(t *testing.T) {
t.Parallel()
tests := []struct {
name string
chars int
baseChars int
want int
}{
{name: "empty output minimum", chars: 0, baseChars: 1000, want: 1},
{name: "single char", chars: 1, baseChars: 1000, want: 1},
{name: "exactly one base", chars: 1000, baseChars: 1000, want: 1},
{name: "one over base", chars: 1001, baseChars: 1000, want: 2},
{name: "default base", chars: 2000, baseChars: 0, want: 2},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
if got := calculateAIPointCost(tt.chars, tt.baseChars); got != tt.want {
t.Fatalf("calculateAIPointCost(%d, %d) = %d, want %d", tt.chars, tt.baseChars, got, tt.want)
}
})
}
}
func TestComplianceAIPointCharsTrimWhitespace(t *testing.T) {
t.Parallel()
got := countAIPointChars(" " + strings.Repeat("字", 1001) + "\n")
if got != 1001 {
t.Fatalf("countAIPointChars = %d, want 1001", got)
}
}
@@ -52,7 +52,7 @@ type AIPointUsageListParams struct {
type AIPointUsageRepository interface {
Create(ctx context.Context, input CreateAIPointUsageLogInput) (int64, error)
MarkCompleted(ctx context.Context, tenantID, id int64, model *string) error
MarkCompleted(ctx context.Context, tenantID, id int64, requestChars, baseChars, points int, model *string) error
MarkRefunded(ctx context.Context, tenantID, id int64, errorMessage string) error
FindPendingByResource(ctx context.Context, tenantID int64, resourceType string, resourceID *int64, resourceUID *string) (*AIPointUsageLog, error)
List(ctx context.Context, params AIPointUsageListParams) ([]AIPointUsageLog, error)
@@ -97,18 +97,21 @@ RETURNING id
return id, err
}
func (r *aiPointUsageRepository) MarkCompleted(ctx context.Context, tenantID, id int64, model *string) error {
func (r *aiPointUsageRepository) MarkCompleted(ctx context.Context, tenantID, id int64, requestChars, baseChars, points int, model *string) error {
_, err := r.db.Exec(ctx, `
UPDATE ai_point_usage_logs
SET status = 'completed',
model = $1,
request_chars = $2,
base_chars = $3,
points = $4,
error_message = NULL,
completed_at = COALESCE(completed_at, NOW()),
updated_at = NOW()
WHERE id = $2
AND tenant_id = $3
WHERE id = $5
AND tenant_id = $6
AND status = 'pending'
`, model, id, tenantID)
`, model, requestChars, baseChars, points, id, tenantID)
return err
}
@@ -93,7 +93,7 @@ func (w *KolAssistWorker) Process(ctx context.Context, job tenantapp.KolAssistJo
}
switch task.Status {
case "completed":
w.completeAIPoints(ctx, job, "")
w.completeAIPoints(ctx, job, "", string(task.ResultJSON))
return nil
case "failed":
errorMessage := "kol assist task failed"
@@ -262,14 +262,14 @@ func (w *KolAssistWorker) processTask(ctx context.Context, job tenantapp.KolAssi
if err := w.repo.MarkCompleted(ctx, job.TenantID, job.TaskID, resultJSON); err != nil {
return fmt.Errorf("mark kol assist completed: %w", err)
}
w.completeAIPoints(ctx, job, result.Model)
w.completeAIPoints(ctx, job, result.Model, result.Content)
return nil
}
func (w *KolAssistWorker) completeAIPoints(ctx context.Context, job tenantapp.KolAssistJob, model string) {
func (w *KolAssistWorker) completeAIPoints(ctx context.Context, job tenantapp.KolAssistJob, model string, meteredOutputText string) {
resourceType := "kol_assist_task"
resourceUID := job.TaskID
if err := tenantapp.CompletePendingAIPointsByResource(ctx, w.pool, job.TenantID, resourceType, nil, &resourceUID, model); err != nil && w.logger != nil {
if _, err := tenantapp.CompletePendingAIPointsByResource(ctx, w.pool, w.cache, job.TenantID, resourceType, nil, &resourceUID, model, meteredOutputText); err != nil && w.logger != nil {
w.logger.Warn("kol assist ai points confirmation failed",
zap.Error(err),
zap.String("task_id", job.TaskID),