add some scriptsmain

1 files changed, 109 insertions, 4 deletions

diff --git a/internal/httpapi/api.go b/internal/httpapi/api.go
index ed7939f..becfbf3 100644
--- a/internal/httpapi/api.go
+++ b/internal/httpapi/api.go
@@ -10,13 +10,16 @@ import (
 	"io"
 	"log/slog"
 	"net/http"
+	"runtime/debug"
 	"strings"
 	"time"
 
 	"aigw/internal/apierror"
 	"aigw/internal/auth"
+	"aigw/internal/billing"
 	"aigw/internal/catalog"
 	"aigw/internal/domain"
+	"aigw/internal/limits"
 	"aigw/internal/provider"
 	"aigw/internal/routing"
 	"aigw/internal/telemetry"
@@ -31,6 +34,11 @@ type API struct {
 	router        *routing.Router
 	forwarder     *provider.Forwarder
 	usageSink     telemetry.UsageSink
+	billingMeter  billing.Meter
+	limiter       *limits.Limiter
+	usageRecorder interface {
+		RecordUsage(context.Context, domain.UsageEvent) error
+	}
 	metrics       *telemetry.Metrics
 	logger        *slog.Logger
 	maxBodyBytes  int64
@@ -43,6 +51,11 @@ type Options struct {
 	Router        *routing.Router
 	Forwarder     *provider.Forwarder
 	UsageSink     telemetry.UsageSink
+	BillingMeter  billing.Meter
+	Limiter       *limits.Limiter
+	UsageRecorder interface {
+		RecordUsage(context.Context, domain.UsageEvent) error
+	}
 	Metrics       *telemetry.Metrics
 	Logger        *slog.Logger
 	MaxBodyBytes  int64
@@ -56,6 +69,9 @@ func New(options Options) *API {
 		router:        options.Router,
 		forwarder:     options.Forwarder,
 		usageSink:     options.UsageSink,
+		billingMeter:  options.BillingMeter,
+		limiter:       options.Limiter,
+		usageRecorder: options.UsageRecorder,
 		metrics:       options.Metrics,
 		logger:        options.Logger,
 		maxBodyBytes:  options.MaxBodyBytes,
@@ -134,6 +150,25 @@ func (a *API) serveInference(w http.ResponseWriter, r *http.Request, protocol do
 		apierror.Write(w, apierror.Error{Status: http.StatusBadRequest, Type: "invalid_params", Message: "Parameter model is required and the body must be valid JSON"}, requestID)
 		return
 	}
+	if a.limiter != nil {
+		lease, limitErr := a.limiter.Acquire(r.Context(), principal, body)
+		if limitErr != nil {
+			w.Header().Set("Retry-After", "60")
+			typeName, message := "rate_limit_exceeded", "Project request rate limit exceeded"
+			switch {
+			case errors.Is(limitErr, limits.ErrTokensExceeded):
+				typeName, message = "token_rate_limit_exceeded", "Project token rate limit exceeded"
+			case errors.Is(limitErr, limits.ErrConcurrencyLimit):
+				typeName, message = "concurrency_limit_exceeded", "Project concurrent request limit exceeded"
+			}
+			apierror.Write(w, apierror.Error{Status: http.StatusTooManyRequests, Type: typeName, Message: message}, requestID)
+			a.recordUsageOnly(r, domain.UsageEvent{RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID,
+				PublicModel: envelope.Model, Protocol: protocol, Stream: envelope.Stream, StatusCode: http.StatusTooManyRequests, Success: false,
+				ErrorType: typeName, StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds()})
+			return
+		}
+		defer lease.Release()
+	}
 
 	routes, err := a.router.Plan(envelope.Model, protocol)
 	if err != nil {
@@ -144,6 +179,41 @@ func (a *API) serveInference(w http.ResponseWriter, r *http.Request, protocol do
 		}
 		return
 	}
+	if a.billingMeter != nil {
+		model, modelErr := a.catalog.Model(envelope.Model)
+		if modelErr != nil {
+			apierror.Write(w, apierror.Error{Status: http.StatusNotFound, Type: "invalid_model", Message: "Model does not exist"}, requestID)
+			return
+		}
+		policy := domain.LimitPolicy{}
+		if a.limiter != nil {
+			policy, _ = a.limiter.Policy(principal.ProjectID)
+		}
+		if err := a.billingMeter.Authorize(r.Context(), billing.Authorization{
+			RequestID: requestID, Principal: principal, Model: model, Body: body, Policy: policy,
+		}); err != nil {
+			if errors.Is(err, billing.ErrInsufficientBalance) {
+				apierror.Write(w, apierror.Error{Status: http.StatusPaymentRequired, Type: "insufficient_balance", Message: "Account balance is insufficient"}, requestID)
+				a.recordUsageOnly(r, domain.UsageEvent{RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID,
+					PublicModel: envelope.Model, Protocol: protocol, Stream: envelope.Stream, StatusCode: http.StatusPaymentRequired, Success: false,
+					ErrorType: "insufficient_balance", StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds()})
+				return
+			}
+			if errors.Is(err, billing.ErrQuotaExceeded) {
+				apierror.Write(w, apierror.Error{Status: http.StatusTooManyRequests, Type: "monthly_quota_exceeded", Message: "Project monthly spend quota exceeded"}, requestID)
+				a.recordUsageOnly(r, domain.UsageEvent{RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID,
+					PublicModel: envelope.Model, Protocol: protocol, Stream: envelope.Stream, StatusCode: http.StatusTooManyRequests, Success: false,
+					ErrorType: "monthly_quota_exceeded", StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds()})
+				return
+			}
+			a.logger.Error("billing_authorization_failed", "request_id", requestID, "tenant_id", principal.TenantID, "error", err)
+			apierror.Write(w, apierror.Error{Status: http.StatusServiceUnavailable, Type: "billing_unavailable", Message: "Billing service is temporarily unavailable"}, requestID)
+			a.recordUsageOnly(r, domain.UsageEvent{RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID,
+				PublicModel: envelope.Model, Protocol: protocol, Stream: envelope.Stream, StatusCode: http.StatusServiceUnavailable, Success: false,
+				ErrorType: "billing_unavailable", StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds()})
+			return
+		}
+	}
 
 	result, err := a.forwarder.Forward(r.Context(), protocol, requestID, body, r.Header, routes)
 	if err != nil {
@@ -155,7 +225,7 @@ func (a *API) serveInference(w http.ResponseWriter, r *http.Request, protocol do
 		} else {
 			apierror.Write(w, apierror.Error{Status: status, Type: errorType, Message: "No upstream provider is currently available"}, requestID)
 		}
-		a.publishUsage(domain.UsageEvent{
+		a.finishUsage(r, domain.UsageEvent{
 			RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID,
 			PublicModel: envelope.Model, Protocol: protocol, Stream: envelope.Stream, StatusCode: status,
 			Success: false, ErrorType: errorType, Attempts: result.Attempts, StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds(),
@@ -167,7 +237,7 @@ func (a *API) serveInference(w http.ResponseWriter, r *http.Request, protocol do
 	if result.Response.StatusCode >= 400 {
 		gatewayError := normalizeProviderError(result.Response)
 		apierror.Write(w, gatewayError, requestID)
-		a.publishUsage(domain.UsageEvent{
+		a.finishUsage(r, domain.UsageEvent{
 			RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID,
 			PublicModel: envelope.Model, ProviderID: result.Route.Provider.ID, UpstreamModel: result.Route.UpstreamModel,
 			Protocol: protocol, Stream: envelope.Stream, StatusCode: gatewayError.Status, Success: false,
@@ -192,7 +262,7 @@ func (a *API) serveInference(w http.ResponseWriter, r *http.Request, protocol do
 	if copyErr != nil {
 		errorType = "stream_interrupted"
 	}
-	a.publishUsage(domain.UsageEvent{
+	a.finishUsage(r, domain.UsageEvent{
 		RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID,
 		PublicModel: envelope.Model, ProviderID: result.Route.Provider.ID, UpstreamModel: result.Route.UpstreamModel,
 		Protocol: protocol, Stream: stream, StatusCode: result.Response.StatusCode, Success: success,
@@ -254,6 +324,41 @@ func (a *API) publishUsage(event domain.UsageEvent) {
 	}
 }
 
+func (a *API) finishUsage(r *http.Request, event domain.UsageEvent) {
+	settled := false
+	if a.billingMeter != nil {
+		ctx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), 5*time.Second)
+		err := a.billingMeter.Settle(ctx, event)
+		cancel()
+		if err != nil {
+			a.logger.Error("billing_settlement_failed", "request_id", event.RequestID, "tenant_id", event.TenantID, "error", err)
+		} else {
+			settled = true
+		}
+	}
+	if a.usageRecorder != nil && !settled {
+		ctx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), 5*time.Second)
+		err := a.usageRecorder.RecordUsage(ctx, event)
+		cancel()
+		if err != nil {
+			a.logger.Error("usage_persistence_failed", "request_id", event.RequestID, "tenant_id", event.TenantID, "error", err)
+		}
+	}
+	a.publishUsage(event)
+}
+
+func (a *API) recordUsageOnly(r *http.Request, event domain.UsageEvent) {
+	if a.usageRecorder != nil {
+		ctx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), 5*time.Second)
+		err := a.usageRecorder.RecordUsage(ctx, event)
+		cancel()
+		if err != nil {
+			a.logger.Error("usage_persistence_failed", "request_id", event.RequestID, "tenant_id", event.TenantID, "error", err)
+		}
+	}
+	a.publishUsage(event)
+}
+
 func (a *API) withRequestID(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		id := newRequestID()
@@ -266,7 +371,7 @@ func (a *API) recoverPanics(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		defer func() {
 			if recovered := recover(); recovered != nil {
-				a.logger.Error("request_panic", "request_id", requestIDFrom(r.Context()), "error", recovered)
+				a.logger.Error("request_panic", "request_id", requestIDFrom(r.Context()), "error", recovered, "stack", string(debug.Stack()))
 				apierror.Write(w, apierror.Error{Status: http.StatusInternalServerError, Type: "internal_server_error", Message: "Internal server error"}, requestIDFrom(r.Context()))
 			}
 		}()