diff options
Diffstat (limited to '')
| -rw-r--r-- | internal/billing/ledger.go | 171 | ||||
| -rw-r--r-- | internal/billing/service.go | 352 | ||||
| -rw-r--r-- | internal/billing/service_test.go | 190 | ||||
| -rw-r--r-- | internal/billing/stripe.go | 203 | ||||
| -rw-r--r-- | internal/billing/types.go | 83 |
5 files changed, 999 insertions, 0 deletions
diff --git a/internal/billing/ledger.go b/internal/billing/ledger.go new file mode 100644 index 0000000..28cbe4c --- /dev/null +++ b/internal/billing/ledger.go @@ -0,0 +1,171 @@ +package billing + +import ( + "context" + "errors" + "fmt" + "math" + "strings" + + "github.com/jackc/pgx/v5" +) + +func (s *Service) ListAccounts(ctx context.Context, tenantID string) ([]Account, error) { + query := ` + SELECT t.id::text, t.name, COALESCE(w.currency, $1), COALESCE(w.balance_micros, 0), + COALESCE(w.reserved_micros, 0), COALESCE(w.balance_micros - w.reserved_micros, 0), + COALESCE(w.updated_at, t.created_at) + FROM tenants t LEFT JOIN tenant_wallets w ON w.tenant_id = t.id` + args := []any{s.currency} + if strings.TrimSpace(tenantID) != "" { + query += ` WHERE t.id=$2` + args = append(args, tenantID) + } + query += ` ORDER BY t.name` + rows, err := s.db.Query(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("query billing accounts: %w", err) + } + defer rows.Close() + result := make([]Account, 0) + for rows.Next() { + var item Account + if err := rows.Scan(&item.TenantID, &item.TenantName, &item.Currency, &item.BalanceMicros, + &item.ReservedMicros, &item.AvailableMicros, &item.UpdatedAt); err != nil { + return nil, fmt.Errorf("scan billing account: %w", err) + } + result = append(result, item) + } + return result, rows.Err() +} + +func (s *Service) ListLedger(ctx context.Context, tenantID string, limit int) ([]LedgerEntry, error) { + if limit < 1 || limit > 500 { + limit = 200 + } + query := ` + SELECT id::text, tenant_id::text, COALESCE(project_id::text, ''), currency, amount_micros, + balance_after_micros, kind, source_type, source_id, description, created_at + FROM billing_ledger` + args := []any{} + if strings.TrimSpace(tenantID) != "" { + query += ` WHERE tenant_id = $1 ORDER BY created_at DESC LIMIT $2` + args = append(args, tenantID, limit) + } else { + query += ` ORDER BY created_at DESC LIMIT $1` + args = append(args, limit) + } + rows, err := s.db.Query(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("query billing ledger: %w", err) + } + defer rows.Close() + result := make([]LedgerEntry, 0) + for rows.Next() { + var item LedgerEntry + if err := rows.Scan(&item.ID, &item.TenantID, &item.ProjectID, &item.Currency, &item.AmountMicros, + &item.BalanceAfterMicros, &item.Kind, &item.SourceType, &item.SourceID, + &item.Description, &item.CreatedAt); err != nil { + return nil, fmt.Errorf("scan billing ledger: %w", err) + } + result = append(result, item) + } + return result, rows.Err() +} + +func (s *Service) AdjustBalance(ctx context.Context, input AdjustmentInput) (LedgerEntry, error) { + input.TenantID = strings.TrimSpace(input.TenantID) + input.Description = normalizeDescription(input.Description) + if input.TenantID == "" || input.AmountMicros == 0 { + return LedgerEntry{}, ErrInvalidAmount + } + tx, err := s.db.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted}) + if err != nil { + return LedgerEntry{}, fmt.Errorf("begin balance adjustment: %w", err) + } + defer tx.Rollback(ctx) + if _, err := tx.Exec(ctx, ` + INSERT INTO tenant_wallets (tenant_id, currency) VALUES ($1, $2) + ON CONFLICT (tenant_id) DO NOTHING`, input.TenantID, s.currency); err != nil { + return LedgerEntry{}, fmt.Errorf("ensure adjustment wallet: %w", err) + } + var currency string + var balance, held int64 + if err := tx.QueryRow(ctx, `SELECT currency, balance_micros, reserved_micros FROM tenant_wallets WHERE tenant_id = $1 FOR UPDATE`, input.TenantID).Scan(¤cy, &balance, &held); err != nil { + return LedgerEntry{}, fmt.Errorf("lock adjustment wallet: %w", err) + } + if currency != s.currency { + return LedgerEntry{}, fmt.Errorf("tenant wallet currency %s does not match %s", currency, s.currency) + } + if input.AmountMicros > 0 && balance > math.MaxInt64-input.AmountMicros { + return LedgerEntry{}, ErrInvalidAmount + } + newBalance := balance + input.AmountMicros + if newBalance < held || newBalance < 0 { + return LedgerEntry{}, ErrInsufficientBalance + } + if _, err := tx.Exec(ctx, `UPDATE tenant_wallets SET balance_micros = $2, updated_at = now() WHERE tenant_id = $1`, input.TenantID, newBalance); err != nil { + return LedgerEntry{}, fmt.Errorf("apply balance adjustment: %w", err) + } + sourceID := "adj_" + randomHex(16) + var result LedgerEntry + err = tx.QueryRow(ctx, ` + INSERT INTO billing_ledger (tenant_id, currency, amount_micros, balance_after_micros, kind, source_type, source_id, description) + VALUES ($1,$2,$3,$4,'adjustment','admin',$5,$6) + RETURNING id::text, tenant_id::text, '', currency, amount_micros, balance_after_micros, + kind, source_type, source_id, description, created_at`, + input.TenantID, s.currency, input.AmountMicros, newBalance, sourceID, input.Description, + ).Scan(&result.ID, &result.TenantID, &result.ProjectID, &result.Currency, &result.AmountMicros, + &result.BalanceAfterMicros, &result.Kind, &result.SourceType, &result.SourceID, + &result.Description, &result.CreatedAt) + if err != nil { + return LedgerEntry{}, fmt.Errorf("write adjustment ledger entry: %w", err) + } + if err := tx.Commit(ctx); err != nil { + return LedgerEntry{}, fmt.Errorf("commit balance adjustment: %w", err) + } + return result, nil +} + +func (s *Service) createTopUpOrder(ctx context.Context, input CheckoutInput) (string, int64, error) { + if !s.stripeEnabled { + return "", 0, ErrStripeDisabled + } + if strings.TrimSpace(input.TenantID) == "" || input.AmountMinor < s.minTopUpMinor || input.AmountMinor > s.maxTopUpMinor { + return "", 0, ErrInvalidAmount + } + amountMicros, err := minorToMicros(s.currency, input.AmountMinor) + if err != nil { + return "", 0, err + } + var orderID string + err = s.db.QueryRow(ctx, ` + INSERT INTO topup_orders (tenant_id, amount_minor, amount_micros, currency) + VALUES ($1,$2,$3,$4) RETURNING id::text`, input.TenantID, input.AmountMinor, amountMicros, s.currency, + ).Scan(&orderID) + if err != nil { + return "", 0, fmt.Errorf("create top-up order: %w", err) + } + return orderID, amountMicros, nil +} + +func minorToMicros(currency string, amount int64) (int64, error) { + if amount <= 0 { + return 0, ErrInvalidAmount + } + factor := int64(10_000) + switch strings.ToLower(currency) { + case "bif", "clp", "djf", "gnf", "jpy", "kmf", "krw", "mga", "pyg", "rwf", "ugx", "vnd", "vuv", "xaf", "xof", "xpf": + factor = 1_000_000 + case "bhd", "jod", "kwd", "omr", "tnd": + factor = 1_000 + } + if amount > math.MaxInt64/factor { + return 0, ErrInvalidAmount + } + return amount * factor, nil +} + +func isNotFound(err error) bool { + return errors.Is(err, pgx.ErrNoRows) +} diff --git a/internal/billing/service.go b/internal/billing/service.go new file mode 100644 index 0000000..b87a0bd --- /dev/null +++ b/internal/billing/service.go @@ -0,0 +1,352 @@ +package billing + +import ( + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "math/big" + "strings" + "time" + + "aigw/internal/domain" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +const microsPerUnit = int64(1_000_000) + +type Service struct { + db *pgxpool.Pool + currency string + defaultMaxOutputTokens int64 + minTopUpMinor int64 + maxTopUpMinor int64 + stripeEnabled bool + stripeWebhookSecret string + stripeSuccessURL string + stripeCancelURL string + integrationIdentifier string + createStripeCheckout stripeCheckoutCreator +} + +func New(ctx context.Context, options Options) (*Service, error) { + db, err := pgxpool.New(ctx, options.DatabaseURL) + if err != nil { + return nil, fmt.Errorf("configure billing PostgreSQL: %w", err) + } + if err := db.Ping(ctx); err != nil { + db.Close() + return nil, fmt.Errorf("connect billing PostgreSQL: %w", err) + } + service := &Service{ + db: db, currency: options.Currency, defaultMaxOutputTokens: options.DefaultMaxOutputTokens, + minTopUpMinor: options.MinTopUpMinor, maxTopUpMinor: options.MaxTopUpMinor, + stripeEnabled: options.StripeEnabled, stripeWebhookSecret: options.StripeWebhookSecret, + stripeSuccessURL: options.StripeSuccessURL, stripeCancelURL: options.StripeCancelURL, + integrationIdentifier: "aigw_balance_" + randomLetters(8), + } + if options.StripeEnabled { + service.createStripeCheckout = newStripeCheckoutCreator(options.StripeAPIKey) + } + return service, nil +} + +func (s *Service) Close() { + s.db.Close() +} + +func (s *Service) StripeEnabled() bool { + return s.stripeEnabled +} + +func (s *Service) Currency() string { + return s.currency +} + +func (s *Service) Authorize(ctx context.Context, input Authorization) error { + if input.RequestID == "" || input.Principal.TenantID == "" || input.Principal.ProjectID == "" || input.Principal.KeyID == "" { + return errors.New("billing authorization identity is incomplete") + } + reserved, err := reservationCost(input.Model, input.Body, s.defaultMaxOutputTokens) + if err != nil { + return err + } + tx, err := s.db.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted}) + if err != nil { + return fmt.Errorf("begin billing authorization: %w", err) + } + defer tx.Rollback(ctx) + if _, err := tx.Exec(ctx, ` + INSERT INTO tenant_wallets (tenant_id, currency) VALUES ($1, $2) + ON CONFLICT (tenant_id) DO NOTHING`, input.Principal.TenantID, s.currency); err != nil { + return fmt.Errorf("ensure tenant wallet: %w", err) + } + var currency string + var balance, held int64 + if err := tx.QueryRow(ctx, ` + SELECT currency, balance_micros, reserved_micros FROM tenant_wallets + WHERE tenant_id = $1 FOR UPDATE`, input.Principal.TenantID).Scan(¤cy, &balance, &held); err != nil { + return fmt.Errorf("lock tenant wallet: %w", err) + } + if currency != s.currency { + return fmt.Errorf("tenant wallet currency %s does not match billing currency %s", currency, s.currency) + } + if input.Policy.MonthlySpendMicros > 0 { + period := time.Date(time.Now().UTC().Year(), time.Now().UTC().Month(), 1, 0, 0, 0, 0, time.UTC) + nextPeriod := period.AddDate(0, 1, 0) + var used, pending int64 + if err := tx.QueryRow(ctx, `SELECT + COALESCE((SELECT cost_micros FROM usage_monthly_rollups WHERE project_id=$1 AND period_start=$2),0), + COALESCE((SELECT sum(reserved_micros) FROM billing_reservations WHERE project_id=$1 AND status='pending' AND created_at >= $2 AND created_at < $3),0)`, + input.Principal.ProjectID, period, nextPeriod).Scan(&used, &pending); err != nil { + return fmt.Errorf("read monthly spend quota: %w", err) + } + if reserved > input.Policy.MonthlySpendMicros || used > input.Policy.MonthlySpendMicros-reserved || pending > input.Policy.MonthlySpendMicros-used-reserved { + return ErrQuotaExceeded + } + } + if balance-held < reserved { + return ErrInsufficientBalance + } + if _, err := tx.Exec(ctx, ` + INSERT INTO billing_reservations ( + request_id, tenant_id, project_id, key_id, public_model, currency, reserved_micros, + input_price_micros_per_million, output_price_micros_per_million, + cache_read_price_micros_per_million, cache_write_price_micros_per_million) + VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11)`, + input.RequestID, input.Principal.TenantID, input.Principal.ProjectID, input.Principal.KeyID, + input.Model.ID, s.currency, reserved, input.Model.InputPriceMicrosPerMillion, + input.Model.OutputPriceMicrosPerMillion, input.Model.CacheReadPriceMicrosPerMillion, + input.Model.CacheWritePriceMicrosPerMillion); err != nil { + return fmt.Errorf("create billing reservation: %w", err) + } + if _, err := tx.Exec(ctx, ` + UPDATE tenant_wallets SET reserved_micros = reserved_micros + $2, updated_at = now() + WHERE tenant_id = $1`, input.Principal.TenantID, reserved); err != nil { + return fmt.Errorf("reserve tenant balance: %w", err) + } + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("commit billing authorization: %w", err) + } + return nil +} + +func (s *Service) Settle(ctx context.Context, event domain.UsageEvent) error { + tx, err := s.db.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted}) + if err != nil { + return fmt.Errorf("begin usage settlement: %w", err) + } + defer tx.Rollback(ctx) + + var tenantID, projectID, keyID, modelID, currency, status string + var reserved, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64 + err = tx.QueryRow(ctx, ` + SELECT tenant_id::text, project_id::text, key_id::text, public_model, currency, reserved_micros, status, + input_price_micros_per_million, output_price_micros_per_million, + cache_read_price_micros_per_million, cache_write_price_micros_per_million + FROM billing_reservations WHERE request_id = $1 FOR UPDATE`, event.RequestID, + ).Scan(&tenantID, &projectID, &keyID, &modelID, ¤cy, &reserved, &status, + &inputPrice, &outputPrice, &cacheReadPrice, &cacheWritePrice) + if err != nil { + return fmt.Errorf("load billing reservation: %w", err) + } + if status != "pending" { + return tx.Commit(ctx) + } + + actualCost := int64(0) + if event.StatusCode >= 200 && event.StatusCode < 300 { + actualCost, err = usageCost(event.Usage, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice) + if err != nil { + return err + } + } + var balance, held int64 + if err := tx.QueryRow(ctx, `SELECT balance_micros, reserved_micros FROM tenant_wallets WHERE tenant_id = $1 FOR UPDATE`, tenantID).Scan(&balance, &held); err != nil { + return fmt.Errorf("lock wallet for settlement: %w", err) + } + charged, err := collectibleCharge(actualCost, balance, held, reserved) + if err != nil { + return err + } + uncollected := actualCost - charged + newBalance := balance - charged + if _, err := tx.Exec(ctx, ` + UPDATE tenant_wallets SET balance_micros = $2, reserved_micros = reserved_micros - $3, updated_at = now() + WHERE tenant_id = $1`, tenantID, newBalance, reserved); err != nil { + return fmt.Errorf("settle tenant wallet: %w", err) + } + reservationStatus := "released" + if actualCost > 0 { + reservationStatus = "settled" + } + if _, err := tx.Exec(ctx, ` + UPDATE billing_reservations + SET status = $2, actual_cost_micros = $3, charged_micros = $4, uncollected_micros = $5, settled_at = now() + WHERE request_id = $1`, event.RequestID, reservationStatus, actualCost, charged, uncollected); err != nil { + return fmt.Errorf("update billing reservation: %w", err) + } + usageAlreadyRecorded := false + if err := tx.QueryRow(ctx, `SELECT true FROM usage_events WHERE request_id=$1 FOR UPDATE`, event.RequestID).Scan(&usageAlreadyRecorded); err != nil && !errors.Is(err, pgx.ErrNoRows) { + return fmt.Errorf("lock existing usage event: %w", err) + } + if usageAlreadyRecorded { + if _, err := tx.Exec(ctx, `UPDATE usage_events SET cost_micros=$2, charged_micros=$3, uncollected_micros=$4 WHERE request_id=$1`, + event.RequestID, actualCost, charged, uncollected); err != nil { + return fmt.Errorf("apply usage charge: %w", err) + } + } else if _, err := tx.Exec(ctx, ` + INSERT INTO usage_events ( + request_id, tenant_id, project_id, key_id, public_model, provider_id, upstream_model, + protocol, stream, status_code, success, error_type, attempts, started_at, duration_ms, + input_tokens, output_tokens, total_tokens, cache_creation_input_tokens, cache_read_input_tokens, + cost_micros, charged_micros, uncollected_micros) + VALUES ($1,$2,$3,$4,$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23) + ON CONFLICT (request_id) DO NOTHING`, + event.RequestID, tenantID, projectID, keyID, modelID, event.ProviderID, event.UpstreamModel, + string(event.Protocol), event.Stream, event.StatusCode, event.Success, event.ErrorType, event.Attempts, + event.StartedAt, event.DurationMS, event.Usage.InputTokens, event.Usage.OutputTokens, + event.Usage.TotalTokens, event.Usage.CacheCreationInputTokens, event.Usage.CacheReadInputTokens, + actualCost, charged, uncollected); err != nil { + return fmt.Errorf("persist usage event: %w", err) + } + if charged > 0 { + if _, err := tx.Exec(ctx, ` + INSERT INTO billing_ledger (tenant_id, project_id, currency, amount_micros, balance_after_micros, kind, source_type, source_id, description) + VALUES ($1,$2,$3,$4,$5,'usage','request',$6,$7) + ON CONFLICT (source_type, source_id) DO NOTHING`, tenantID, projectID, currency, -charged, newBalance, event.RequestID, modelID); err != nil { + return fmt.Errorf("write usage ledger entry: %w", err) + } + } + period := time.Date(event.StartedAt.UTC().Year(), event.StartedAt.UTC().Month(), 1, 0, 0, 0, 0, time.UTC) + if usageAlreadyRecorded { + if _, err := tx.Exec(ctx, `UPDATE usage_monthly_rollups SET cost_micros=cost_micros+$3, + charged_micros=charged_micros+$4, uncollected_micros=uncollected_micros+$5, updated_at=now() + WHERE project_id=$1 AND period_start=$2`, projectID, period, actualCost, charged, uncollected); err != nil { + return fmt.Errorf("apply usage rollup charge: %w", err) + } + } else if _, err := tx.Exec(ctx, ` + INSERT INTO usage_monthly_rollups + (period_start, tenant_id, project_id, request_count, successful_requests, input_tokens, output_tokens, total_tokens, cost_micros, charged_micros, uncollected_micros) + VALUES ($1,$2,$3,1,$4,$5,$6,$7,$8,$9,$10) + ON CONFLICT (project_id, period_start) DO UPDATE SET request_count=usage_monthly_rollups.request_count+1, + successful_requests=usage_monthly_rollups.successful_requests+EXCLUDED.successful_requests, + input_tokens=usage_monthly_rollups.input_tokens+EXCLUDED.input_tokens, + output_tokens=usage_monthly_rollups.output_tokens+EXCLUDED.output_tokens, + total_tokens=usage_monthly_rollups.total_tokens+EXCLUDED.total_tokens, + cost_micros=usage_monthly_rollups.cost_micros+EXCLUDED.cost_micros, + charged_micros=usage_monthly_rollups.charged_micros+EXCLUDED.charged_micros, + uncollected_micros=usage_monthly_rollups.uncollected_micros+EXCLUDED.uncollected_micros, + updated_at=now()`, period, tenantID, projectID, boolToInt(event.Success), event.Usage.InputTokens, + event.Usage.OutputTokens, event.Usage.TotalTokens, actualCost, charged, uncollected); err != nil { + return fmt.Errorf("update usage monthly rollup: %w", err) + } + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("commit usage settlement: %w", err) + } + return nil +} + +func boolToInt(value bool) int { + if value { + return 1 + } + return 0 +} + +func reservationCost(model domain.Model, body []byte, defaultMaxOutput int64) (int64, error) { + maxOutput := defaultMaxOutput + var limits struct { + MaxTokens int64 `json:"max_tokens"` + MaxCompletionTokens int64 `json:"max_completion_tokens"` + MaxOutputTokens int64 `json:"max_output_tokens"` + } + if json.Unmarshal(body, &limits) == nil { + explicitMax := int64(0) + for _, value := range []int64{limits.MaxTokens, limits.MaxCompletionTokens, limits.MaxOutputTokens} { + if value > explicitMax { + explicitMax = value + } + } + if explicitMax > 0 { + maxOutput = explicitMax + } + } + cacheReservePrice := model.CacheReadPriceMicrosPerMillion + if model.CacheWritePriceMicrosPerMillion > cacheReservePrice { + cacheReservePrice = model.CacheWritePriceMicrosPerMillion + } + return calculateCost(int64(len(body)), maxOutput, 0, int64(len(body)), + model.InputPriceMicrosPerMillion, model.OutputPriceMicrosPerMillion, + 0, cacheReservePrice) +} + +func usageCost(usage domain.Usage, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64) (int64, error) { + return calculateCost(usage.InputTokens, usage.OutputTokens, usage.CacheReadInputTokens, + usage.CacheCreationInputTokens, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice) +} + +func calculateCost(input, output, cacheRead, cacheWrite, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64) (int64, error) { + values := []int64{input, output, cacheRead, cacheWrite, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice} + for _, value := range values { + if value < 0 { + return 0, errors.New("billing values cannot be negative") + } + } + total := new(big.Int) + for _, pair := range [][2]int64{{input, inputPrice}, {output, outputPrice}, {cacheRead, cacheReadPrice}, {cacheWrite, cacheWritePrice}} { + total.Add(total, new(big.Int).Mul(big.NewInt(pair[0]), big.NewInt(pair[1]))) + } + if total.Sign() == 0 { + return 0, nil + } + total.Add(total, big.NewInt(microsPerUnit-1)) + total.Div(total, big.NewInt(microsPerUnit)) + if !total.IsInt64() { + return 0, errors.New("calculated charge exceeds supported range") + } + return total.Int64(), nil +} + +func collectibleCharge(actualCost, balance, held, reservation int64) (int64, error) { + if actualCost < 0 || balance < 0 || held < 0 || reservation < 0 || held > balance || reservation > held { + return 0, errors.New("wallet reservation invariant violated") + } + spendable := balance - (held - reservation) + if actualCost > spendable { + return spendable, nil + } + return actualCost, nil +} + +func randomHex(bytes int) string { + buffer := make([]byte, bytes) + if _, err := rand.Read(buffer); err != nil { + return fmt.Sprintf("%08x", time.Now().UnixNano())[:bytes*2] + } + return hex.EncodeToString(buffer) +} + +func randomLetters(length int) string { + const letters = "abcdefghijklmnopqrstuvwxyz" + buffer := make([]byte, length) + if _, err := rand.Read(buffer); err != nil { + return strings.Repeat("a", length) + } + for i := range buffer { + buffer[i] = letters[int(buffer[i])%len(letters)] + } + return string(buffer) +} + +func normalizeDescription(value string) string { + value = strings.TrimSpace(value) + if len(value) > 240 { + value = value[:240] + } + return value +} diff --git a/internal/billing/service_test.go b/internal/billing/service_test.go new file mode 100644 index 0000000..17f8143 --- /dev/null +++ b/internal/billing/service_test.go @@ -0,0 +1,190 @@ +package billing + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "aigw/internal/controlplane" + "aigw/internal/domain" + + "github.com/stripe/stripe-go/v86" + "github.com/stripe/stripe-go/v86/webhook" +) + +func TestUsageCostUsesFixedPointAndRoundsOnce(t *testing.T) { + usage := domain.Usage{InputTokens: 3, OutputTokens: 2, CacheReadInputTokens: 5} + cost, err := usageCost(usage, 150_000, 600_000, 30_000, 0) + if err != nil { + t.Fatal(err) + } + // (3*150000 + 2*600000 + 5*30000) / 1e6 = 1.8 micros. + if cost != 2 { + t.Fatalf("cost = %d, want 2", cost) + } +} + +func TestReservationUsesExplicitOutputLimit(t *testing.T) { + model := domain.Model{InputPriceMicrosPerMillion: 1_000_000, OutputPriceMicrosPerMillion: 1_000_000} + body := []byte(`{"max_tokens":8192}`) + cost, err := reservationCost(model, body, 4096) + if err != nil { + t.Fatal(err) + } + want := int64(len(body)) + 8192 + if cost != want { + t.Fatalf("reservation = %d, want %d", cost, want) + } +} + +func TestMinorToMicrosSupportsCurrencyExponents(t *testing.T) { + tests := []struct { + currency string + minor int64 + want int64 + }{{"usd", 123, 1_230_000}, {"jpy", 123, 123_000_000}, {"bhd", 123, 123_000}} + for _, test := range tests { + got, err := minorToMicros(test.currency, test.minor) + if err != nil { + t.Fatalf("%s: %v", test.currency, err) + } + if got != test.want { + t.Fatalf("%s: got %d, want %d", test.currency, got, test.want) + } + } +} + +func TestCollectibleChargePreservesOtherReservations(t *testing.T) { + charged, err := collectibleCharge(100, 100, 80, 40) + if err != nil { + t.Fatal(err) + } + if charged != 60 { + t.Fatalf("charged = %d, want 60", charged) + } + if remainingBalance, remainingHeld := int64(100)-charged, int64(80)-40; remainingBalance < remainingHeld { + t.Fatalf("remaining balance %d does not cover held %d", remainingBalance, remainingHeld) + } +} + +func TestIntegrationIdentifierSuffixUsesLetters(t *testing.T) { + value := randomLetters(8) + if len(value) != 8 { + t.Fatalf("length = %d", len(value)) + } + for _, char := range value { + if char < 'a' || char > 'z' { + t.Fatalf("non-letter suffix %q", value) + } + } +} + +func TestWebhookRejectsInvalidSignatureBeforeProcessing(t *testing.T) { + service := &Service{stripeWebhookSecret: "whsec_test"} + request := httptest.NewRequest(http.MethodPost, "/billing/stripe/webhook", strings.NewReader(`{"id":"evt_fake"}`)) + request.Header.Set("Stripe-Signature", "invalid") + response := httptest.NewRecorder() + service.WebhookHandler().ServeHTTP(response, request) + if response.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400", response.Code) + } +} + +func TestStripeWebhookCreditsPaidOrderExactlyOncePostgres(t *testing.T) { + databaseURL := os.Getenv("AIGW_TEST_DATABASE_URL") + if databaseURL == "" { + t.Skip("AIGW_TEST_DATABASE_URL is not set") + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := controlplane.MigrateDatabase(ctx, databaseURL); err != nil { + t.Fatal(err) + } + service, err := New(ctx, Options{ + DatabaseURL: databaseURL, Currency: "usd", MinTopUpMinor: 500, MaxTopUpMinor: 1_000_000, + StripeEnabled: true, StripeAPIKey: "rk_test_placeholder", StripeWebhookSecret: "whsec_integration_test", + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(service.Close) + + var tenantID string + slug := fmt.Sprintf("stripe-%d", time.Now().UnixNano()) + if err := service.db.QueryRow(ctx, `INSERT INTO tenants (slug, name) VALUES ($1,'Stripe integration') RETURNING id::text`, slug).Scan(&tenantID); err != nil { + t.Fatal(err) + } + eventID := fmt.Sprintf("evt_aigw_%d", time.Now().UnixNano()) + t.Cleanup(func() { + cleanupCtx := context.Background() + for _, statement := range []struct { + query string + arg string + }{ + {`DELETE FROM stripe_webhook_events WHERE event_id=$1`, eventID}, + {`DELETE FROM billing_ledger WHERE tenant_id=$1`, tenantID}, + {`DELETE FROM tenant_wallets WHERE tenant_id=$1`, tenantID}, + {`DELETE FROM topup_orders WHERE tenant_id=$1`, tenantID}, + {`DELETE FROM tenants WHERE id=$1`, tenantID}, + } { + if _, cleanupErr := service.db.Exec(cleanupCtx, statement.query, statement.arg); cleanupErr != nil { + t.Errorf("cleanup Stripe integration data: %v", cleanupErr) + } + } + }) + const amountMinor int64 = 2500 + orderID, amountMicros, err := service.createTopUpOrder(ctx, CheckoutInput{TenantID: tenantID, AmountMinor: amountMinor}) + if err != nil { + t.Fatal(err) + } + sessionID := fmt.Sprintf("cs_test_aigw_%d", time.Now().UnixNano()) + payload, err := json.Marshal(map[string]any{ + "id": eventID, "object": "event", "api_version": stripe.APIVersion, + "type": string(stripe.EventTypeCheckoutSessionCompleted), + "data": map[string]any{"object": map[string]any{ + "id": sessionID, "object": "checkout.session", "client_reference_id": orderID, + "amount_total": amountMinor, "currency": "usd", "payment_status": "paid", + }}, + }) + if err != nil { + t.Fatal(err) + } + signed := webhook.GenerateTestSignedPayload(&webhook.UnsignedPayload{Payload: payload, Secret: service.stripeWebhookSecret}) + for delivery := 0; delivery < 2; delivery++ { + request := httptest.NewRequest(http.MethodPost, "/billing/stripe/webhook", strings.NewReader(string(payload))) + request.Header.Set("Stripe-Signature", signed.Header) + response := httptest.NewRecorder() + service.WebhookHandler().ServeHTTP(response, request) + if response.Code != http.StatusOK { + t.Fatalf("delivery %d status = %d, body = %s", delivery+1, response.Code, response.Body.String()) + } + } + + var balance int64 + if err := service.db.QueryRow(ctx, `SELECT balance_micros FROM tenant_wallets WHERE tenant_id=$1`, tenantID).Scan(&balance); err != nil { + t.Fatal(err) + } + if balance != amountMicros { + t.Fatalf("balance = %d, want %d", balance, amountMicros) + } + var ledgerCount, webhookCount int + var orderStatus string + if err := service.db.QueryRow(ctx, `SELECT count(*) FROM billing_ledger WHERE source_type='stripe_checkout' AND source_id=$1`, sessionID).Scan(&ledgerCount); err != nil { + t.Fatal(err) + } + if err := service.db.QueryRow(ctx, `SELECT count(*) FROM stripe_webhook_events WHERE event_id=$1`, eventID).Scan(&webhookCount); err != nil { + t.Fatal(err) + } + if err := service.db.QueryRow(ctx, `SELECT status FROM topup_orders WHERE id=$1`, orderID).Scan(&orderStatus); err != nil { + t.Fatal(err) + } + if ledgerCount != 1 || webhookCount != 1 || orderStatus != "paid" { + t.Fatalf("ledger=%d webhook=%d order=%s, want 1/1/paid", ledgerCount, webhookCount, orderStatus) + } +} diff --git a/internal/billing/stripe.go b/internal/billing/stripe.go new file mode 100644 index 0000000..b06eb6a --- /dev/null +++ b/internal/billing/stripe.go @@ -0,0 +1,203 @@ +package billing + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + + "github.com/jackc/pgx/v5" + "github.com/stripe/stripe-go/v86" + "github.com/stripe/stripe-go/v86/webhook" +) + +const maxWebhookBodyBytes = 1 << 20 + +type stripeCheckoutCreator func(context.Context, *stripe.CheckoutSessionCreateParams) (*stripe.CheckoutSession, error) + +func newStripeCheckoutCreator(apiKey string) stripeCheckoutCreator { + client := stripe.NewClient(apiKey) + return client.V1CheckoutSessions.Create +} + +func (s *Service) CreateCheckout(ctx context.Context, input CheckoutInput) (CheckoutResult, error) { + orderID, _, err := s.createTopUpOrder(ctx, input) + if err != nil { + return CheckoutResult{}, err + } + params := &stripe.CheckoutSessionCreateParams{ + Mode: stripe.String("payment"), + ClientReferenceID: stripe.String(orderID), + IntegrationIdentifier: stripe.String(s.integrationIdentifier), + SuccessURL: stripe.String(s.stripeSuccessURL), + CancelURL: stripe.String(s.stripeCancelURL), + Metadata: map[string]string{ + "aigw_topup_order_id": orderID, + "aigw_tenant_id": strings.TrimSpace(input.TenantID), + }, + LineItems: []*stripe.CheckoutSessionCreateLineItemParams{{ + Quantity: stripe.Int64(1), + PriceData: &stripe.CheckoutSessionCreateLineItemPriceDataParams{ + Currency: stripe.String(s.currency), + UnitAmount: stripe.Int64(input.AmountMinor), + ProductData: &stripe.CheckoutSessionCreateLineItemPriceDataProductDataParams{ + Name: stripe.String("AIGW prepaid balance"), + Description: stripe.String("Prepaid API usage credit"), + }, + }, + }}, + } + params.SetIdempotencyKey("aigw_topup_" + orderID) + session, err := s.createStripeCheckout(ctx, params) + if err != nil { + _, _ = s.db.Exec(ctx, `UPDATE topup_orders SET status = 'failed' WHERE id = $1 AND status = 'pending'`, orderID) + return CheckoutResult{}, fmt.Errorf("create Stripe Checkout Session: %w", err) + } + if session.ID == "" || session.URL == "" { + return CheckoutResult{}, errors.New("Stripe returned an incomplete Checkout Session") + } + if _, err := s.db.Exec(ctx, ` + UPDATE topup_orders SET stripe_session_id = $2, checkout_url = $3 + WHERE id = $1`, orderID, session.ID, session.URL); err != nil { + return CheckoutResult{}, fmt.Errorf("persist Stripe Checkout Session: %w", err) + } + return CheckoutResult{OrderID: orderID, SessionID: session.ID, URL: session.URL}, nil +} + +func (s *Service) WebhookHandler() http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + w.Header().Set("Allow", http.MethodPost) + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + body, err := io.ReadAll(io.LimitReader(r.Body, maxWebhookBodyBytes+1)) + if err != nil || len(body) > maxWebhookBodyBytes { + http.Error(w, "invalid webhook body", http.StatusBadRequest) + return + } + event, err := webhook.ConstructEvent(body, r.Header.Get("Stripe-Signature"), s.stripeWebhookSecret) + if err != nil { + http.Error(w, "invalid webhook signature", http.StatusBadRequest) + return + } + if err := s.processStripeEvent(r.Context(), event); err != nil { + if errors.Is(err, ErrInvalidAmount) || isNotFound(err) { + http.Error(w, "invalid checkout event", http.StatusBadRequest) + return + } + http.Error(w, "webhook processing failed", http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"received":true}`+"\n") + }) +} + +func (s *Service) processStripeEvent(ctx context.Context, event stripe.Event) error { + typeName := string(event.Type) + switch event.Type { + case stripe.EventTypeCheckoutSessionCompleted, + stripe.EventTypeCheckoutSessionAsyncPaymentSucceeded, + stripe.EventTypeCheckoutSessionAsyncPaymentFailed, + stripe.EventTypeCheckoutSessionExpired: + default: + return nil + } + if event.Data == nil { + return ErrInvalidAmount + } + var session stripe.CheckoutSession + if err := json.Unmarshal(event.Data.Raw, &session); err != nil { + return ErrInvalidAmount + } + if event.ID == "" || session.ID == "" || session.ClientReferenceID == "" { + return ErrInvalidAmount + } + tx, err := s.db.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted}) + if err != nil { + return err + } + defer tx.Rollback(ctx) + tag, err := tx.Exec(ctx, ` + INSERT INTO stripe_webhook_events (event_id, event_type) VALUES ($1,$2) + ON CONFLICT (event_id) DO NOTHING`, event.ID, typeName) + if err != nil { + return fmt.Errorf("record Stripe event: %w", err) + } + if tag.RowsAffected() == 0 { + return tx.Commit(ctx) + } + + var tenantID, currency, status string + var amountMinor, amountMicros int64 + var storedSessionID *string + err = tx.QueryRow(ctx, ` + SELECT tenant_id::text, amount_minor, amount_micros, currency, status, stripe_session_id + FROM topup_orders WHERE id = $1 FOR UPDATE`, session.ClientReferenceID, + ).Scan(&tenantID, &amountMinor, &amountMicros, ¤cy, &status, &storedSessionID) + if err != nil { + return err + } + if (storedSessionID != nil && *storedSessionID != session.ID) || amountMinor != session.AmountTotal || currency != string(session.Currency) { + return ErrInvalidAmount + } + if event.Type == stripe.EventTypeCheckoutSessionAsyncPaymentFailed || event.Type == stripe.EventTypeCheckoutSessionExpired { + orderStatus := "failed" + if event.Type == stripe.EventTypeCheckoutSessionExpired { + orderStatus = "expired" + } + if _, err := tx.Exec(ctx, `UPDATE topup_orders SET status = $2, stripe_session_id = COALESCE(stripe_session_id, $3) WHERE id = $1 AND status = 'pending'`, session.ClientReferenceID, orderStatus, session.ID); err != nil { + return err + } + _, err = tx.Exec(ctx, `UPDATE stripe_webhook_events SET processed_at = now() WHERE event_id = $1`, event.ID) + if err != nil { + return err + } + return tx.Commit(ctx) + } + if session.PaymentStatus != stripe.CheckoutSessionPaymentStatusPaid { + _, err = tx.Exec(ctx, `UPDATE stripe_webhook_events SET processed_at = now() WHERE event_id = $1`, event.ID) + if err != nil { + return err + } + return tx.Commit(ctx) + } + if status != "paid" { + if _, err := tx.Exec(ctx, ` + INSERT INTO tenant_wallets (tenant_id, currency) VALUES ($1,$2) + ON CONFLICT (tenant_id) DO NOTHING`, tenantID, currency); err != nil { + return err + } + var walletCurrency string + var balance int64 + if err := tx.QueryRow(ctx, `SELECT currency, balance_micros FROM tenant_wallets WHERE tenant_id = $1 FOR UPDATE`, tenantID).Scan(&walletCurrency, &balance); err != nil { + return err + } + if walletCurrency != currency || amountMicros <= 0 || balance > int64(^uint64(0)>>1)-amountMicros { + return ErrInvalidAmount + } + newBalance := balance + amountMicros + if _, err := tx.Exec(ctx, `UPDATE tenant_wallets SET balance_micros = $2, updated_at = now() WHERE tenant_id = $1`, tenantID, newBalance); err != nil { + return err + } + if _, err := tx.Exec(ctx, ` + INSERT INTO billing_ledger (tenant_id, currency, amount_micros, balance_after_micros, kind, source_type, source_id, description) + VALUES ($1,$2,$3,$4,'topup','stripe_checkout',$5,'Stripe balance top-up') + ON CONFLICT (source_type, source_id) DO NOTHING`, tenantID, currency, amountMicros, newBalance, session.ID); err != nil { + return err + } + if _, err := tx.Exec(ctx, ` + UPDATE topup_orders SET status = 'paid', stripe_session_id = COALESCE(stripe_session_id, $2), paid_at = now() + WHERE id = $1`, session.ClientReferenceID, session.ID); err != nil { + return err + } + } + if _, err := tx.Exec(ctx, `UPDATE stripe_webhook_events SET processed_at = now() WHERE event_id = $1`, event.ID); err != nil { + return err + } + return tx.Commit(ctx) +} diff --git a/internal/billing/types.go b/internal/billing/types.go new file mode 100644 index 0000000..24694cc --- /dev/null +++ b/internal/billing/types.go @@ -0,0 +1,83 @@ +package billing + +import ( + "context" + "errors" + "time" + + "aigw/internal/domain" +) + +var ( + ErrInsufficientBalance = errors.New("insufficient balance") + ErrStripeDisabled = errors.New("Stripe top-ups are disabled") + ErrInvalidAmount = errors.New("invalid amount") + ErrQuotaExceeded = errors.New("monthly spend quota exceeded") +) + +type Meter interface { + Authorize(context.Context, Authorization) error + Settle(context.Context, domain.UsageEvent) error +} + +type Authorization struct { + RequestID string + Principal domain.Principal + Model domain.Model + Body []byte + Policy domain.LimitPolicy +} + +type Options struct { + DatabaseURL string + Currency string + DefaultMaxOutputTokens int64 + MinTopUpMinor int64 + MaxTopUpMinor int64 + StripeEnabled bool + StripeAPIKey string + StripeWebhookSecret string + StripeSuccessURL string + StripeCancelURL string +} + +type Account struct { + TenantID string `json:"tenant_id"` + TenantName string `json:"tenant_name"` + Currency string `json:"currency"` + BalanceMicros int64 `json:"balance_micros"` + ReservedMicros int64 `json:"reserved_micros"` + AvailableMicros int64 `json:"available_micros"` + UpdatedAt time.Time `json:"updated_at"` +} + +type LedgerEntry struct { + ID string `json:"id"` + TenantID string `json:"tenant_id"` + ProjectID string `json:"project_id,omitempty"` + Currency string `json:"currency"` + AmountMicros int64 `json:"amount_micros"` + BalanceAfterMicros int64 `json:"balance_after_micros"` + Kind string `json:"kind"` + SourceType string `json:"source_type"` + SourceID string `json:"source_id"` + Description string `json:"description"` + CreatedAt time.Time `json:"created_at"` +} + +type AdjustmentInput struct { + TenantID string `json:"tenant_id"` + AmountMicros int64 `json:"amount_micros"` + Description string `json:"description"` +} + +type CheckoutInput struct { + TenantID string `json:"tenant_id"` + AmountMinor int64 `json:"amount_minor"` +} + +type CheckoutResult struct { + OrderID string `json:"order_id"` + SessionID string `json:"session_id"` + URL string `json:"url"` +} |