mirrors/pikasTech-HWLAB
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add v03 user billing entitlements

Browse Source
This commit is contained in:
lyon
2026-06-14 15:34:13 +08:00
parent ed80e6dd1b
commit c84ee72eb9
12 changed files with 712 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/userbilling/migrations/0002_billing_entitlements.sql
+78
View File
@@ -0,0 +1,78 @@
CREATE TABLE IF NOT EXISTS hwlab_billing_plans (
id TEXT PRIMARY KEY,
display_name TEXT NOT NULL DEFAULT '',
status TEXT NOT NULL DEFAULT 'active',
description TEXT NOT NULL DEFAULT '',
metadata JSONB NOT NULL DEFAULT '{}'::jsonb,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
CHECK (status IN ('active', 'disabled'))
);
INSERT INTO hwlab_billing_plans (id, display_name, status, description, metadata)
VALUES ('default', 'Default', 'active', 'Default HWLAB v0.3 resource plan', '{"source":"hwlab-v03","limitSemantics":"zero-means-unlimited"}'::jsonb)
ON CONFLICT (id) DO UPDATE SET
display_name = EXCLUDED.display_name,
status = EXCLUDED.status,
description = EXCLUDED.description,
metadata = hwlab_billing_plans.metadata || EXCLUDED.metadata,
updated_at = now();
CREATE TABLE IF NOT EXISTS hwlab_resource_entitlements (
id TEXT PRIMARY KEY,
plan_id TEXT NOT NULL REFERENCES hwlab_billing_plans(id) ON DELETE CASCADE,
resource_type TEXT NOT NULL,
enabled BOOLEAN NOT NULL DEFAULT TRUE,
monthly_quota_credits BIGINT NOT NULL DEFAULT 0,
max_concurrent_reservations INTEGER NOT NULL DEFAULT 0,
rpm_limit INTEGER NOT NULL DEFAULT 0,
metadata JSONB NOT NULL DEFAULT '{}'::jsonb,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
UNIQUE (plan_id, resource_type),
CHECK (resource_type <> ''),
CHECK (monthly_quota_credits >= 0),
CHECK (max_concurrent_reservations >= 0),
CHECK (rpm_limit >= 0)
);
CREATE INDEX IF NOT EXISTS hwlab_resource_entitlements_plan_idx ON hwlab_resource_entitlements(plan_id);
CREATE INDEX IF NOT EXISTS hwlab_resource_entitlements_resource_idx ON hwlab_resource_entitlements(resource_type);
INSERT INTO hwlab_resource_entitlements (id, plan_id, resource_type, enabled, monthly_quota_credits, max_concurrent_reservations, rpm_limit, metadata)
VALUES
('ent_default_code_agent', 'default', 'code_agent', TRUE, 0, 0, 0, '{"serviceIds":["hwlab-code-agent"],"limitSemantics":"zero-means-unlimited"}'::jsonb),
('ent_default_aipod', 'default', 'aipod', TRUE, 0, 0, 0, '{"serviceIds":["hwlab-aipod"],"limitSemantics":"zero-means-unlimited"}'::jsonb),
('ent_default_hwpod', 'default', 'hwpod', TRUE, 0, 0, 0, '{"serviceIds":["hwlab-hwpod"],"limitSemantics":"zero-means-unlimited"}'::jsonb)
ON CONFLICT (plan_id, resource_type) DO UPDATE SET
enabled = EXCLUDED.enabled,
monthly_quota_credits = EXCLUDED.monthly_quota_credits,
max_concurrent_reservations = EXCLUDED.max_concurrent_reservations,
rpm_limit = EXCLUDED.rpm_limit,
metadata = hwlab_resource_entitlements.metadata || EXCLUDED.metadata,
updated_at = now();
ALTER TABLE hwlab_billing_reservations ADD COLUMN IF NOT EXISTS resource_type TEXT NOT NULL DEFAULT 'generic';
ALTER TABLE hwlab_usage_records ADD COLUMN IF NOT EXISTS resource_type TEXT NOT NULL DEFAULT 'generic';
UPDATE hwlab_billing_reservations
SET resource_type = CASE
WHEN lower(service_id) LIKE '%code-agent%' OR lower(service_id) LIKE '%code_agent%' THEN 'code_agent'
WHEN lower(service_id) LIKE '%aipod%' THEN 'aipod'
WHEN lower(service_id) LIKE '%hwpod%' THEN 'hwpod'
ELSE resource_type
END
WHERE resource_type = 'generic';
UPDATE hwlab_usage_records
SET resource_type = CASE
WHEN lower(service_id) LIKE '%code-agent%' OR lower(service_id) LIKE '%code_agent%' THEN 'code_agent'
WHEN lower(service_id) LIKE '%aipod%' THEN 'aipod'
WHEN lower(service_id) LIKE '%hwpod%' THEN 'hwpod'
ELSE resource_type
END
WHERE resource_type = 'generic';
CREATE INDEX IF NOT EXISTS hwlab_billing_reservations_user_resource_idx ON hwlab_billing_reservations(user_id, resource_type, status, expires_at);
CREATE INDEX IF NOT EXISTS hwlab_billing_reservations_user_resource_created_idx ON hwlab_billing_reservations(user_id, resource_type, created_at);
CREATE INDEX IF NOT EXISTS hwlab_usage_records_user_resource_created_idx ON hwlab_usage_records(user_id, resource_type, created_at);
internal/userbilling/service.go
+384 -34
View File
@@ -79,11 +79,12 @@ type Principal struct {
}
type serviceUsageSummary struct {
ServiceID string `json:"serviceId"`
Credits int64 `json:"credits"`
Quantity int64 `json:"quantity"`
RecordCount int64 `json:"recordCount"`
LastUsedAt time.Time `json:"lastUsedAt"`
ServiceID string `json:"serviceId"`
ResourceType string `json:"resourceType"`
Credits int64 `json:"credits"`
Quantity int64 `json:"quantity"`
RecordCount int64 `json:"recordCount"`
LastUsedAt time.Time `json:"lastUsedAt"`
}
type ledgerSummaryRow struct {
@@ -98,6 +99,7 @@ type ledgerSummaryRow struct {
type reservationSummaryRow struct {
ReservationID string `json:"reservationId"`
ServiceID string `json:"serviceId"`
ResourceType string `json:"resourceType"`
EstimatedCredits int64 `json:"estimatedCredits"`
EstimatedTokens int64 `json:"estimatedTokens"`
Status string `json:"status"`
@@ -105,13 +107,42 @@ type reservationSummaryRow struct {
CreatedAt time.Time `json:"createdAt"`
}
type billingPlanSummary struct {
ID string `json:"id"`
DisplayName string `json:"displayName"`
Status string `json:"status"`
Description string `json:"description"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type resourceEntitlementSummary struct {
ID string `json:"id,omitempty"`
PlanID string `json:"planId,omitempty"`
ResourceType string `json:"resourceType"`
Enabled bool `json:"enabled"`
MonthlyQuotaCredits int64 `json:"monthlyQuotaCredits"`
MonthlyQuotaUnlimited bool `json:"monthlyQuotaUnlimited"`
MonthlyUsageCredits int64 `json:"monthlyUsageCredits"`
MonthlyRemainingCredits int64 `json:"monthlyRemainingCredits"`
MaxConcurrentReservations int64 `json:"maxConcurrentReservations"`
ConcurrencyUnlimited bool `json:"concurrencyUnlimited"`
ActiveReservations int64 `json:"activeReservations"`
RPMLimit int64 `json:"rpmLimit"`
RPMUnlimited bool `json:"rpmUnlimited"`
RPMUsed int64 `json:"rpmUsed"`
LimitSemantics string `json:"limitSemantics"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type adminBillingUserSummary struct {
User User `json:"user"`
PlanID string `json:"planId"`
Credits creditSummary `json:"credits"`
Usage userUsageAggregate `json:"usage"`
APIKeys apiKeyAggregate `json:"apiKeys"`
Reservations reservationAdminSummary `json:"reservations"`
User User `json:"user"`
PlanID string `json:"planId"`
Plan billingPlanSummary `json:"plan,omitempty"`
Entitlements []resourceEntitlementSummary `json:"entitlements,omitempty"`
Credits creditSummary `json:"credits"`
Usage userUsageAggregate `json:"usage"`
APIKeys apiKeyAggregate `json:"apiKeys"`
Reservations reservationAdminSummary `json:"reservations"`
}
type creditSummary struct {
@@ -160,6 +191,11 @@ type adminUserDetail struct {
Reservations []reservationSummaryRow `json:"reservations"`
}
type billingPlanDetail struct {
Plan billingPlanSummary `json:"plan"`
Entitlements []resourceEntitlementSummary `json:"entitlements"`
}
type reservationAdminSummary struct {
ActiveCount int64 `json:"activeCount"`
Active []reservationSummaryRow `json:"active"`
@@ -286,6 +322,7 @@ func (s *Server) routes() {
s.route(http.MethodPost, "/internal/billing/record", s.handleBillingRecord)
s.route(http.MethodPost, "/internal/billing/release", s.handleBillingRelease)
s.route(http.MethodGet, "/internal/admin/billing/summary", s.handleAdminBillingSummary)
s.route(http.MethodGet, "/internal/admin/billing/plans", s.handleAdminBillingPlans)
s.route(http.MethodPost, "/internal/admin/credits/adjust", s.handleAdminCreditAdjust)
s.mux.HandleFunc("/internal/admin/users", s.handleAdminUsers)
s.mux.HandleFunc("/internal/admin/users/", s.handleAdminUserByID)
@@ -663,6 +700,12 @@ func (s *Server) handleBillingSummary(w http.ResponseWriter, r *http.Request) {
writeAPIError(w, http.StatusInternalServerError, "billing_summary_failed", "could not load active reservations")
return
}
planID, plan, entitlements, err := s.accountPlanContext(ctx, principal.UserID)
if err != nil {
s.logDatabaseError("billing_summary_entitlements", err)
writeAPIError(w, http.StatusInternalServerError, "billing_summary_failed", "could not load billing plan entitlements")
return
}
writeJSON(w, http.StatusOK, map[string]any{
"contractVersion": "user-billing-summary-v1",
@@ -674,6 +717,9 @@ func (s *Server) handleBillingSummary(w http.ResponseWriter, r *http.Request) {
"reserved": reserved,
"available": balance - reserved,
},
"planId": planID,
"plan": plan,
"entitlements": entitlements,
"usage": map[string]any{
"totalCredits": totalCredits,
"totalQuantity": totalQuantity,
@@ -903,6 +949,11 @@ func (s *Server) handleBillingPreflight(w http.ResponseWriter, r *http.Request)
writeJSON(w, http.StatusPaymentRequired, map[string]any{"allowed": false, "reason": "insufficient_credits"})
return
}
var limitErr billingLimitError
if errors.As(err, &limitErr) {
writeJSON(w, limitErr.Status, map[string]any{"allowed": false, "reason": limitErr.Code, "error": apiError{Code: limitErr.Code, Message: limitErr.Message}})
return
}
writeAPIError(w, http.StatusInternalServerError, "billing_preflight_failed", "could not reserve credits")
return
}
@@ -1114,6 +1165,7 @@ func (s *Server) handleAdminUpdateUser(w http.ResponseWriter, r *http.Request, u
Password *string `json:"password"`
Role *string `json:"role"`
Status *string `json:"status"`
PlanID *string `json:"planId"`
}
if !decodeJSON(w, r, &req) {
return
@@ -1175,7 +1227,15 @@ func (s *Server) handleAdminUpdateUser(w http.ResponseWriter, r *http.Request, u
}
ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
defer cancel()
if len(updates) == 0 {
planID := ""
if req.PlanID != nil {
planID = strings.TrimSpace(*req.PlanID)
if planID == "" {
writeAPIError(w, http.StatusBadRequest, "invalid_plan", "planId is required when updating a billing plan")
return
}
}
if len(updates) == 0 && planID == "" {
detail, err := s.adminUserDetail(ctx, userID)
if errors.Is(err, sql.ErrNoRows) {
writeAPIError(w, http.StatusNotFound, "user_not_found", "user not found")
@@ -1188,6 +1248,19 @@ func (s *Server) handleAdminUpdateUser(w http.ResponseWriter, r *http.Request, u
writeJSON(w, http.StatusOK, map[string]any{"contractVersion": "user-billing-admin-user-v1", "updated": false, "detail": detail, "valuesRedacted": true})
return
}
if len(updates) == 0 {
detail, err := s.updateUserPlan(ctx, userID, planID)
if errors.Is(err, sql.ErrNoRows) {
writeAPIError(w, http.StatusNotFound, "user_or_plan_not_found", "user or plan not found")
return
}
if err != nil {
writeAPIError(w, http.StatusBadRequest, "plan_update_failed", err.Error())
return
}
writeJSON(w, http.StatusOK, map[string]any{"contractVersion": "user-billing-admin-user-v1", "updated": true, "detail": detail, "valuesRedacted": true})
return
}
query := `UPDATE hwlab_users SET ` + strings.Join(updates, ", ") + `, updated_at = now() WHERE id = $1 RETURNING id, email, username, display_name, status, role, email_verified, created_at`
var user User
err := s.db.QueryRowContext(ctx, query, args...).Scan(&user.ID, &user.Email, &user.Username, &user.DisplayName, &user.Status, &user.Role, &user.EmailVerified, &user.CreatedAt)
@@ -1203,6 +1276,12 @@ func (s *Server) handleAdminUpdateUser(w http.ResponseWriter, r *http.Request, u
writeAPIError(w, http.StatusInternalServerError, "user_update_failed", "could not update user")
return
}
if planID != "" {
if _, err := s.updateUserPlan(ctx, userID, planID); err != nil {
writeAPIError(w, http.StatusBadRequest, "plan_update_failed", err.Error())
return
}
}
writeJSON(w, http.StatusOK, map[string]any{"contractVersion": "user-billing-admin-user-v1", "updated": true, "user": user, "valuesRedacted": true})
}
@@ -1302,8 +1381,31 @@ func (s *Server) handleAdminBillingSummary(w http.ResponseWriter, r *http.Reques
})
}
func (s *Server) handleAdminBillingPlans(w http.ResponseWriter, r *http.Request) {
if !s.internalAllowed(w, r) || !s.databaseAvailable(w) {
return
}
ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
defer cancel()
plans, err := s.billingPlans(ctx)
if err != nil {
s.logDatabaseError("admin_billing_plans", err)
writeAPIError(w, http.StatusInternalServerError, "admin_billing_plans_failed", "could not load billing plans")
return
}
writeJSON(w, http.StatusOK, map[string]any{"contractVersion": "user-billing-admin-plans-v1", "plans": plans, "count": len(plans), "limitSemantics": "zero-means-unlimited", "valuesRedacted": true})
}
var errInsufficientCredits = errors.New("insufficient credits")
type billingLimitError struct {
Code string
Message string
Status int
}
func (err billingLimitError) Error() string { return err.Message }
func (s *Server) reserveCredits(ctx context.Context, userID, service string, credits, tokens int64, idempotencyKey string, metadata map[string]any) (map[string]any, error) {
tx, err := s.db.BeginTx(ctx, &sql.TxOptions{})
if err != nil {
@@ -1313,16 +1415,23 @@ func (s *Server) reserveCredits(ctx context.Context, userID, service string, cre
if idempotencyKey != "" {
var existingID, status string
var estimated int64
err := tx.QueryRowContext(ctx, `SELECT id, status, estimated_credits FROM hwlab_billing_reservations WHERE idempotency_key = $1`, idempotencyKey).Scan(&existingID, &status, &estimated)
var resourceType string
err := tx.QueryRowContext(ctx, `SELECT id, status, estimated_credits, resource_type FROM hwlab_billing_reservations WHERE idempotency_key = $1`, idempotencyKey).Scan(&existingID, &status, &estimated, &resourceType)
if err == nil {
return map[string]any{"allowed": status == "reserved", "reservationId": existingID, "estimatedCredits": estimated, "status": status, "idempotent": true}, tx.Commit()
return map[string]any{"allowed": status == "reserved", "reservationId": existingID, "estimatedCredits": estimated, "resourceType": resourceType, "status": status, "idempotent": true}, tx.Commit()
}
if !errors.Is(err, sql.ErrNoRows) {
return nil, err
}
}
var balance, reserved int64
if err := tx.QueryRowContext(ctx, `SELECT balance_credits, reserved_credits FROM hwlab_credit_accounts WHERE user_id = $1 FOR UPDATE`, userID).Scan(&balance, &reserved); err != nil {
var planID string
if err := tx.QueryRowContext(ctx, `SELECT balance_credits, reserved_credits, plan_id FROM hwlab_credit_accounts WHERE user_id = $1 FOR UPDATE`, userID).Scan(&balance, &reserved, &planID); err != nil {
return nil, err
}
resourceType := resourceTypeForService(service)
entitlement, err := s.evaluateResourceEntitlement(ctx, tx, userID, planID, resourceType, credits)
if err != nil {
return nil, err
}
available := balance - reserved
@@ -1332,7 +1441,7 @@ func (s *Server) reserveCredits(ctx context.Context, userID, service string, cre
reservationID := newID("res")
metadataJSON := jsonObject(metadata)
expiresAt := time.Now().UTC().Add(s.config.ReservationTTL)
_, err = tx.ExecContext(ctx, `INSERT INTO hwlab_billing_reservations (id, user_id, service_id, estimated_credits, estimated_tokens, idempotency_key, metadata, expires_at) VALUES ($1, $2, $3, $4, $5, nullif($6, ''), $7::jsonb, $8)`, reservationID, userID, service, credits, nonNegative(tokens), idempotencyKey, metadataJSON, expiresAt)
_, err = tx.ExecContext(ctx, `INSERT INTO hwlab_billing_reservations (id, user_id, service_id, resource_type, estimated_credits, estimated_tokens, idempotency_key, metadata, expires_at) VALUES ($1, $2, $3, $4, $5, $6, nullif($7, ''), $8::jsonb, $9)`, reservationID, userID, service, resourceType, credits, nonNegative(tokens), idempotencyKey, metadataJSON, expiresAt)
if err != nil {
return nil, err
}
@@ -1343,7 +1452,7 @@ func (s *Server) reserveCredits(ctx context.Context, userID, service string, cre
if err := tx.Commit(); err != nil {
return nil, err
}
return map[string]any{"allowed": true, "reservationId": reservationID, "estimatedCredits": credits, "estimatedTokens": tokens, "availableBefore": available, "expiresAt": expiresAt}, nil
return map[string]any{"allowed": true, "reservationId": reservationID, "estimatedCredits": credits, "estimatedTokens": tokens, "resourceType": resourceType, "planId": planID, "entitlement": entitlement, "availableBefore": available, "expiresAt": expiresAt}, nil
}
func (s *Server) recordUsage(ctx context.Context, r *http.Request, reservationID, token, apiKey, userID, serviceID string, usedCredits, usedTokens int64, idempotencyKey string, metadata map[string]any) (map[string]any, error) {
@@ -1363,9 +1472,10 @@ func (s *Server) recordUsage(ctx context.Context, r *http.Request, reservationID
}
}
var reserved int64
resourceType := ""
if reservationID != "" {
var status string
err := tx.QueryRowContext(ctx, `SELECT user_id, service_id, estimated_credits, status FROM hwlab_billing_reservations WHERE id = $1 FOR UPDATE`, reservationID).Scan(&userID, &serviceID, &reserved, &status)
err := tx.QueryRowContext(ctx, `SELECT user_id, service_id, resource_type, estimated_credits, status FROM hwlab_billing_reservations WHERE id = $1 FOR UPDATE`, reservationID).Scan(&userID, &serviceID, &resourceType, &reserved, &status)
if err != nil {
return nil, err
}
@@ -1379,15 +1489,25 @@ func (s *Server) recordUsage(ctx context.Context, r *http.Request, reservationID
}
userID = principal.UserID
serviceID = strings.TrimSpace(serviceID)
resourceType = resourceTypeForService(serviceID)
}
if serviceID == "" {
return nil, errors.New("serviceId is required")
}
if resourceType == "" {
resourceType = resourceTypeForService(serviceID)
}
credits := s.estimateCredits(usedCredits, usedTokens)
var balance, accountReserved int64
if err := tx.QueryRowContext(ctx, `SELECT balance_credits, reserved_credits FROM hwlab_credit_accounts WHERE user_id = $1 FOR UPDATE`, userID).Scan(&balance, &accountReserved); err != nil {
var planID string
if err := tx.QueryRowContext(ctx, `SELECT balance_credits, reserved_credits, plan_id FROM hwlab_credit_accounts WHERE user_id = $1 FOR UPDATE`, userID).Scan(&balance, &accountReserved, &planID); err != nil {
return nil, err
}
if reservationID == "" {
if _, err := s.evaluateResourceEntitlement(ctx, tx, userID, planID, resourceType, credits); err != nil {
return nil, err
}
}
if reservationID == "" && balance-accountReserved < credits {
return nil, errInsufficientCredits
}
@@ -1405,7 +1525,7 @@ func (s *Server) recordUsage(ctx context.Context, r *http.Request, reservationID
}
recordID := newID("use")
metadataJSON := jsonObject(metadata)
_, err = tx.ExecContext(ctx, `INSERT INTO hwlab_usage_records (id, reservation_id, user_id, service_id, quantity, credits, idempotency_key, metadata) VALUES ($1, nullif($2, ''), $3, $4, $5, $6, nullif($7, ''), $8::jsonb)`, recordID, reservationID, userID, serviceID, nonNegative(usedTokens), credits, idempotencyKey, metadataJSON)
_, err = tx.ExecContext(ctx, `INSERT INTO hwlab_usage_records (id, reservation_id, user_id, service_id, resource_type, quantity, credits, idempotency_key, metadata) VALUES ($1, nullif($2, ''), $3, $4, $5, $6, $7, nullif($8, ''), $9::jsonb)`, recordID, reservationID, userID, serviceID, resourceType, nonNegative(usedTokens), credits, idempotencyKey, metadataJSON)
if err != nil {
return nil, err
}
@@ -1426,7 +1546,7 @@ func (s *Server) recordUsage(ctx context.Context, r *http.Request, reservationID
if err := tx.Commit(); err != nil {
return nil, err
}
return map[string]any{"recordId": recordID, "userId": userID, "serviceId": serviceID, "credits": credits, "balance": newBalance}, nil
return map[string]any{"recordId": recordID, "userId": userID, "serviceId": serviceID, "resourceType": resourceType, "credits": credits, "balance": newBalance}, nil
}
func (s *Server) releaseReservation(ctx context.Context, reservationID, serviceID, reason, idempotencyKey string, metadata map[string]any) (map[string]any, error) {
@@ -1439,9 +1559,9 @@ func (s *Server) releaseReservation(ctx context.Context, reservationID, serviceI
return nil, err
}
defer tx.Rollback()
var userID, actualServiceID, status string
var userID, actualServiceID, resourceType, status string
var reserved int64
err = tx.QueryRowContext(ctx, `SELECT user_id, service_id, estimated_credits, status FROM hwlab_billing_reservations WHERE id = $1 FOR UPDATE`, reservationID).Scan(&userID, &actualServiceID, &reserved, &status)
err = tx.QueryRowContext(ctx, `SELECT user_id, service_id, resource_type, estimated_credits, status FROM hwlab_billing_reservations WHERE id = $1 FOR UPDATE`, reservationID).Scan(&userID, &actualServiceID, &resourceType, &reserved, &status)
if err != nil {
return nil, err
}
@@ -1450,7 +1570,7 @@ func (s *Server) releaseReservation(ctx context.Context, reservationID, serviceI
return nil, fmt.Errorf("reservation service is %s", actualServiceID)
}
if status == "cancelled" || status == "expired" {
return map[string]any{"released": true, "reservationId": reservationID, "userId": userID, "serviceId": actualServiceID, "releasedCredits": int64(0), "status": status, "idempotent": true}, tx.Commit()
return map[string]any{"released": true, "reservationId": reservationID, "userId": userID, "serviceId": actualServiceID, "resourceType": resourceType, "releasedCredits": int64(0), "status": status, "idempotent": true}, tx.Commit()
}
if status != "reserved" {
return nil, fmt.Errorf("reservation status is %s", status)
@@ -1468,8 +1588,8 @@ func (s *Server) releaseReservation(ctx context.Context, reservationID, serviceI
return nil, err
}
releaseMetadata := map[string]any{
"releaseReason": strings.TrimSpace(reason),
"releaseIdempotency": strings.TrimSpace(idempotencyKey),
"releaseReason": strings.TrimSpace(reason),
"releaseIdempotency": strings.TrimSpace(idempotencyKey),
"releaseValuesRedacted": true,
}
for key, value := range metadata {
@@ -1482,7 +1602,7 @@ func (s *Server) releaseReservation(ctx context.Context, reservationID, serviceI
if err := tx.Commit(); err != nil {
return nil, err
}
return map[string]any{"released": true, "reservationId": reservationID, "userId": userID, "serviceId": actualServiceID, "releasedCredits": reserved, "status": "cancelled"}, nil
return map[string]any{"released": true, "reservationId": reservationID, "userId": userID, "serviceId": actualServiceID, "resourceType": resourceType, "releasedCredits": reserved, "status": "cancelled"}, nil
}
func (s *Server) adjustCredits(ctx context.Context, userID string, delta int64, kind, reason, idempotencyKey string, metadata map[string]any) (int64, error) {
@@ -1520,6 +1640,201 @@ func (s *Server) adjustCredits(ctx context.Context, userID string, delta int64,
return newBalance, tx.Commit()
}
type txQueryer interface {
QueryContext(context.Context, string, ...any) (*sql.Rows, error)
QueryRowContext(context.Context, string, ...any) *sql.Row
}
func (s *Server) evaluateResourceEntitlement(ctx context.Context, tx txQueryer, userID, planID, resourceType string, requestedCredits int64) (resourceEntitlementSummary, error) {
plan, err := s.planByID(ctx, tx, planID)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return resourceEntitlementSummary{}, billingLimitError{Code: "plan_not_found", Message: "billing plan is not active", Status: http.StatusForbidden}
}
return resourceEntitlementSummary{}, err
}
if plan.Status != "active" {
return resourceEntitlementSummary{}, billingLimitError{Code: "plan_disabled", Message: "billing plan is disabled", Status: http.StatusForbidden}
}
entitlements, err := s.resourceEntitlements(ctx, tx, planID, userID)
if err != nil {
return resourceEntitlementSummary{}, err
}
var entitlement resourceEntitlementSummary
found := false
for _, item := range entitlements {
if item.ResourceType == resourceType {
entitlement = item
found = true
break
}
}
if !found || !entitlement.Enabled {
return resourceEntitlementSummary{}, billingLimitError{Code: "resource_not_entitled", Message: "resource is not enabled for this billing plan", Status: http.StatusForbidden}
}
if entitlement.MaxConcurrentReservations > 0 && entitlement.ActiveReservations >= entitlement.MaxConcurrentReservations {
return entitlement, billingLimitError{Code: "resource_concurrency_exceeded", Message: "resource concurrency limit exceeded", Status: http.StatusTooManyRequests}
}
if entitlement.RPMLimit > 0 && entitlement.RPMUsed >= entitlement.RPMLimit {
return entitlement, billingLimitError{Code: "resource_rpm_exceeded", Message: "resource RPM limit exceeded", Status: http.StatusTooManyRequests}
}
if entitlement.MonthlyQuotaCredits > 0 && entitlement.MonthlyUsageCredits+requestedCredits > entitlement.MonthlyQuotaCredits {
return entitlement, billingLimitError{Code: "resource_monthly_quota_exceeded", Message: "resource monthly quota exceeded", Status: http.StatusPaymentRequired}
}
return entitlement, nil
}
func (s *Server) planByID(ctx context.Context, queryer txQueryer, planID string) (billingPlanSummary, error) {
planID = strings.TrimSpace(planID)
if planID == "" {
planID = "default"
}
var plan billingPlanSummary
var metadataRaw string
err := queryer.QueryRowContext(ctx, `SELECT id, display_name, status, description, metadata::text FROM hwlab_billing_plans WHERE id = $1`, planID).Scan(&plan.ID, &plan.DisplayName, &plan.Status, &plan.Description, &metadataRaw)
if err != nil {
return plan, err
}
plan.Metadata = parseJSONMap(metadataRaw)
return plan, nil
}
func (s *Server) billingPlans(ctx context.Context) ([]billingPlanDetail, error) {
rows, err := s.db.QueryContext(ctx, `SELECT id, display_name, status, description, metadata::text FROM hwlab_billing_plans ORDER BY id ASC`)
if err != nil {
return nil, err
}
defer rows.Close()
plans := []billingPlanDetail{}
for rows.Next() {
var detail billingPlanDetail
var metadataRaw string
if err := rows.Scan(&detail.Plan.ID, &detail.Plan.DisplayName, &detail.Plan.Status, &detail.Plan.Description, &metadataRaw); err != nil {
return nil, err
}
detail.Plan.Metadata = parseJSONMap(metadataRaw)
plans = append(plans, detail)
}
if err := rows.Err(); err != nil {
return nil, err
}
for i := range plans {
entitlements, err := s.resourceEntitlements(ctx, s.db, plans[i].Plan.ID, "")
if err != nil {
return nil, err
}
plans[i].Entitlements = entitlements
}
return plans, nil
}
func (s *Server) updateUserPlan(ctx context.Context, userID, planID string) (adminUserDetail, error) {
planID = strings.TrimSpace(planID)
if planID == "" {
return adminUserDetail{}, errors.New("planId is required")
}
if _, err := s.planByID(ctx, s.db, planID); err != nil {
return adminUserDetail{}, err
}
res, err := s.db.ExecContext(ctx, `UPDATE hwlab_credit_accounts SET plan_id = $2, updated_at = now() WHERE user_id = $1`, userID, planID)
if err != nil {
return adminUserDetail{}, err
}
affected, _ := res.RowsAffected()
if affected == 0 {
return adminUserDetail{}, sql.ErrNoRows
}
return s.adminUserDetail(ctx, userID)
}
func (s *Server) resourceEntitlements(ctx context.Context, queryer txQueryer, planID, userID string) ([]resourceEntitlementSummary, error) {
rows, err := queryer.QueryContext(ctx, `
SELECT id, plan_id, resource_type, enabled, monthly_quota_credits, max_concurrent_reservations, rpm_limit, metadata::text
FROM hwlab_resource_entitlements
WHERE plan_id = $1
ORDER BY resource_type ASC`, planID)
if err != nil {
return nil, err
}
defer rows.Close()
items := []resourceEntitlementSummary{}
for rows.Next() {
var item resourceEntitlementSummary
var metadataRaw string
if err := rows.Scan(&item.ID, &item.PlanID, &item.ResourceType, &item.Enabled, &item.MonthlyQuotaCredits, &item.MaxConcurrentReservations, &item.RPMLimit, &metadataRaw); err != nil {
return nil, err
}
item.Metadata = parseJSONMap(metadataRaw)
item.LimitSemantics = "zero-means-unlimited"
item.MonthlyQuotaUnlimited = item.MonthlyQuotaCredits == 0
item.ConcurrencyUnlimited = item.MaxConcurrentReservations == 0
item.RPMUnlimited = item.RPMLimit == 0
items = append(items, item)
}
if err := rows.Err(); err != nil {
return nil, err
}
if strings.TrimSpace(userID) == "" || len(items) == 0 {
return items, nil
}
for i := range items {
if err := s.populateEntitlementUsage(ctx, queryer, userID, &items[i]); err != nil {
return nil, err
}
}
return items, nil
}
func (s *Server) populateEntitlementUsage(ctx context.Context, queryer txQueryer, userID string, item *resourceEntitlementSummary) error {
monthStart := time.Now().UTC().Format("2006-01-02")
if err := queryer.QueryRowContext(ctx, `SELECT COALESCE(SUM(credits), 0) FROM hwlab_usage_records WHERE user_id = $1 AND resource_type = $2 AND created_at >= date_trunc('month', $3::timestamptz)`, userID, item.ResourceType, monthStart).Scan(&item.MonthlyUsageCredits); err != nil {
return err
}
if err := queryer.QueryRowContext(ctx, `SELECT COUNT(*) FROM hwlab_billing_reservations WHERE user_id = $1 AND resource_type = $2 AND status = 'reserved' AND expires_at > now()`, userID, item.ResourceType).Scan(&item.ActiveReservations); err != nil {
return err
}
if err := queryer.QueryRowContext(ctx, `SELECT COUNT(*) FROM hwlab_billing_reservations WHERE user_id = $1 AND resource_type = $2 AND created_at >= now() - interval '1 minute'`, userID, item.ResourceType).Scan(&item.RPMUsed); err != nil {
return err
}
if item.MonthlyQuotaCredits > 0 {
item.MonthlyRemainingCredits = item.MonthlyQuotaCredits - item.MonthlyUsageCredits
if item.MonthlyRemainingCredits < 0 {
item.MonthlyRemainingCredits = 0
}
}
return nil
}
func resourceTypeForService(serviceID string) string {
value := strings.ToLower(strings.TrimSpace(serviceID))
switch {
case strings.Contains(value, "code-agent") || strings.Contains(value, "code_agent") || strings.Contains(value, "codex"):
return "code_agent"
case strings.Contains(value, "aipod") || strings.Contains(value, "ai-pod"):
return "aipod"
case strings.Contains(value, "hwpod") || strings.Contains(value, "hw-pod"):
return "hwpod"
case value != "":
return strings.Map(func(r rune) rune {
if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '_' {
return r
}
if r == '-' || r == '.' || r == '/' {
return '_'
}
return -1
}, value)
default:
return "generic"
}
}
func parseJSONMap(raw string) map[string]any {
result := map[string]any{}
_ = json.Unmarshal([]byte(raw), &result)
return result
}
func (s *Server) principalForBilling(ctx context.Context, r *http.Request, token, apiKey, userID string) (*Principal, error) {
if token == "" && apiKey == "" {
token = bearerToken(r)
@@ -1684,6 +1999,19 @@ func (s *Server) accountBalance(ctx context.Context, userID string) (int64, int6
return balance, reserved, err
}
func (s *Server) accountPlanContext(ctx context.Context, userID string) (string, billingPlanSummary, []resourceEntitlementSummary, error) {
planID := "default"
if err := s.db.QueryRowContext(ctx, `SELECT COALESCE(plan_id, 'default') FROM hwlab_credit_accounts WHERE user_id = $1`, userID).Scan(&planID); err != nil {
return "", billingPlanSummary{}, nil, err
}
plan, err := s.planByID(ctx, s.db, planID)
if err != nil {
return planID, billingPlanSummary{}, nil, err
}
entitlements, err := s.resourceEntitlements(ctx, s.db, planID, userID)
return planID, plan, entitlements, err
}
func (s *Server) usageTotals(ctx context.Context, userID string) (int64, int64, int64, error) {
var credits, quantity, count int64
err := s.db.QueryRowContext(ctx, `SELECT COALESCE(SUM(credits), 0), COALESCE(SUM(quantity), 0), COUNT(*) FROM hwlab_usage_records WHERE user_id = $1`, userID).Scan(&credits, &quantity, &count)
@@ -1691,7 +2019,7 @@ func (s *Server) usageTotals(ctx context.Context, userID string) (int64, int64,
}
func (s *Server) usageByService(ctx context.Context, userID string, limit int) ([]serviceUsageSummary, error) {
rows, err := s.db.QueryContext(ctx, `SELECT service_id, COALESCE(SUM(credits), 0), COALESCE(SUM(quantity), 0), COUNT(*), MAX(created_at) FROM hwlab_usage_records WHERE user_id = $1 GROUP BY service_id ORDER BY MAX(created_at) DESC NULLS LAST, service_id ASC LIMIT $2`, userID, limit)
rows, err := s.db.QueryContext(ctx, `SELECT service_id, resource_type, COALESCE(SUM(credits), 0), COALESCE(SUM(quantity), 0), COUNT(*), MAX(created_at) FROM hwlab_usage_records WHERE user_id = $1 GROUP BY service_id, resource_type ORDER BY MAX(created_at) DESC NULLS LAST, service_id ASC LIMIT $2`, userID, limit)
if err != nil {
return nil, err
}
@@ -1699,7 +2027,7 @@ func (s *Server) usageByService(ctx context.Context, userID string, limit int) (
items := []serviceUsageSummary{}
for rows.Next() {
var item serviceUsageSummary
if err := rows.Scan(&item.ServiceID, &item.Credits, &item.Quantity, &item.RecordCount, &item.LastUsedAt); err != nil {
if err := rows.Scan(&item.ServiceID, &item.ResourceType, &item.Credits, &item.Quantity, &item.RecordCount, &item.LastUsedAt); err != nil {
return nil, err
}
items = append(items, item)
@@ -1725,7 +2053,7 @@ func (s *Server) ledgerRows(ctx context.Context, userID string, limit int) ([]le
}
func (s *Server) activeReservations(ctx context.Context, userID string, limit int) ([]reservationSummaryRow, error) {
rows, err := s.db.QueryContext(ctx, `SELECT id, service_id, estimated_credits, estimated_tokens, status, expires_at, created_at FROM hwlab_billing_reservations WHERE user_id = $1 AND status = 'reserved' AND expires_at > now() ORDER BY expires_at ASC, created_at DESC LIMIT $2`, userID, limit)
rows, err := s.db.QueryContext(ctx, `SELECT id, service_id, resource_type, estimated_credits, estimated_tokens, status, expires_at, created_at FROM hwlab_billing_reservations WHERE user_id = $1 AND status = 'reserved' AND expires_at > now() ORDER BY expires_at ASC, created_at DESC LIMIT $2`, userID, limit)
if err != nil {
return nil, err
}
@@ -1733,7 +2061,7 @@ func (s *Server) activeReservations(ctx context.Context, userID string, limit in
items := []reservationSummaryRow{}
for rows.Next() {
var item reservationSummaryRow
if err := rows.Scan(&item.ReservationID, &item.ServiceID, &item.EstimatedCredits, &item.EstimatedTokens, &item.Status, &item.ExpiresAt, &item.CreatedAt); err != nil {
if err := rows.Scan(&item.ReservationID, &item.ServiceID, &item.ResourceType, &item.EstimatedCredits, &item.EstimatedTokens, &item.Status, &item.ExpiresAt, &item.CreatedAt); err != nil {
return nil, err
}
items = append(items, item)
@@ -1819,6 +2147,9 @@ LIMIT $1`, limit)
for i := range users {
users[i].Reservations.Active = reservations[users[i].User.ID]
}
if err := s.enrichAdminBillingUsers(ctx, users); err != nil {
return nil, err
}
return users, nil
}
@@ -1836,9 +2167,9 @@ func (s *Server) activeReservationsForUsers(ctx context.Context, userIDs []strin
limitParam := "$" + strconv.Itoa(len(args)+1)
args = append(args, perUserLimit)
query := `
SELECT user_id, id, service_id, estimated_credits, estimated_tokens, status, expires_at, created_at
SELECT user_id, id, service_id, resource_type, estimated_credits, estimated_tokens, status, expires_at, created_at
FROM (
SELECT user_id, id, service_id, estimated_credits, estimated_tokens, status, expires_at, created_at,
SELECT user_id, id, service_id, resource_type, estimated_credits, estimated_tokens, status, expires_at, created_at,
row_number() OVER (PARTITION BY user_id ORDER BY expires_at ASC, created_at DESC) AS rn
FROM hwlab_billing_reservations
WHERE user_id IN (` + strings.Join(placeholders, ", ") + `) AND status = 'reserved' AND expires_at > now()
@@ -1853,7 +2184,7 @@ ORDER BY expires_at ASC, created_at DESC`
for rows.Next() {
var userID string
var item reservationSummaryRow
if err := rows.Scan(&userID, &item.ReservationID, &item.ServiceID, &item.EstimatedCredits, &item.EstimatedTokens, &item.Status, &item.ExpiresAt, &item.CreatedAt); err != nil {
if err := rows.Scan(&userID, &item.ReservationID, &item.ServiceID, &item.ResourceType, &item.EstimatedCredits, &item.EstimatedTokens, &item.Status, &item.ExpiresAt, &item.CreatedAt); err != nil {
return nil, err
}
items[userID] = append(items[userID], item)
@@ -2036,9 +2367,28 @@ LIMIT `+limitParam+` OFFSET `+offsetParam, args...)
for i := range users {
users[i].Reservations.Active = reservations[users[i].User.ID]
}
if err := s.enrichAdminBillingUsers(ctx, users); err != nil {
return nil, 0, err
}
return users, total, nil
}
func (s *Server) enrichAdminBillingUsers(ctx context.Context, users []adminBillingUserSummary) error {
for i := range users {
plan, err := s.planByID(ctx, s.db, users[i].PlanID)
if err != nil {
return err
}
entitlements, err := s.resourceEntitlements(ctx, s.db, users[i].PlanID, users[i].User.ID)
if err != nil {
return err
}
users[i].Plan = plan
users[i].Entitlements = entitlements
}
return nil
}
func adminUsersWhere(query adminUsersQuery) (string, []any) {
clauses := []string{}
args := []any{}