feat(tenant): bill AI points by output characters
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:
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user