add some scriptsmain

2 files changed, 192 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()))
 			}
 		}()
diff --git a/internal/httpapi/api_test.go b/internal/httpapi/api_test.go
index 66a344b..014b99b 100644
--- a/internal/httpapi/api_test.go
+++ b/internal/httpapi/api_test.go
@@ -3,6 +3,7 @@ package httpapi
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"encoding/json"
 	"io"
 	"log/slog"
@@ -14,6 +15,7 @@ import (
 	"time"
 
 	"aigw/internal/auth"
+	"aigw/internal/billing"
 	"aigw/internal/catalog"
 	"aigw/internal/config"
 	"aigw/internal/domain"
@@ -26,6 +28,20 @@ type captureUsageSink struct {
 	events chan domain.UsageEvent
 }
 
+type fakeBillingMeter struct {
+	authorizeErr error
+	settled      chan domain.UsageEvent
+}
+
+func (m *fakeBillingMeter) Authorize(context.Context, billing.Authorization) error {
+	return m.authorizeErr
+}
+
+func (m *fakeBillingMeter) Settle(_ context.Context, event domain.UsageEvent) error {
+	m.settled <- event
+	return nil
+}
+
 func (s *captureUsageSink) Publish(event domain.UsageEvent) {
 	s.events <- event
 }
@@ -178,7 +194,73 @@ func TestAnthropicHeadersAndPath(t *testing.T) {
 	}
 }
 
+func TestInsufficientBalanceRejectsBeforeCallingUpstream(t *testing.T) {
+	var calls atomic.Int64
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		calls.Add(1)
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer upstream.Close()
+	meter := &fakeBillingMeter{authorizeErr: billing.ErrInsufficientBalance, settled: make(chan domain.UsageEvent, 1)}
+	gateway, sink := newTestGatewayWithBilling(t, []config.ProviderConfig{{
+		ID: "primary", Protocol: domain.ProtocolOpenAI, BaseURL: upstream.URL + "/v1", APIKey: "secret",
+	}}, []config.RouteConfig{{Provider: "primary", UpstreamModel: "model", Weight: 1}}, meter)
+	defer gateway.Close()
+
+	response := postOpenAI(t, gateway.URL, false)
+	defer response.Body.Close()
+	if response.StatusCode != http.StatusPaymentRequired {
+		t.Fatalf("status = %d, want 402", response.StatusCode)
+	}
+	if calls.Load() != 0 {
+		t.Fatalf("upstream calls = %d, want 0", calls.Load())
+	}
+	select {
+	case <-meter.settled:
+		t.Fatal("rejected request was settled")
+	default:
+	}
+	select {
+	case event := <-sink.events:
+		if event.StatusCode != http.StatusPaymentRequired || event.ErrorType != "insufficient_balance" {
+			t.Fatalf("unexpected rejection usage: %+v", event)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("rejected request did not emit usage")
+	}
+}
+
+func TestSuccessfulRequestIsSettled(t *testing.T) {
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		_, _ = io.WriteString(w, `{"choices":[],"usage":{"prompt_tokens":2,"completion_tokens":3,"total_tokens":5}}`)
+	}))
+	defer upstream.Close()
+	meter := &fakeBillingMeter{settled: make(chan domain.UsageEvent, 1)}
+	gateway, _ := newTestGatewayWithBilling(t, []config.ProviderConfig{{
+		ID: "primary", Protocol: domain.ProtocolOpenAI, BaseURL: upstream.URL + "/v1", APIKey: "secret",
+	}}, []config.RouteConfig{{Provider: "primary", UpstreamModel: "model", Weight: 1}}, meter)
+	defer gateway.Close()
+
+	response := postOpenAI(t, gateway.URL, false)
+	defer response.Body.Close()
+	if response.StatusCode != http.StatusOK {
+		t.Fatalf("status = %d, want 200", response.StatusCode)
+	}
+	select {
+	case event := <-meter.settled:
+		if event.Usage.TotalTokens != 5 {
+			t.Fatalf("settled usage = %+v", event.Usage)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("request was not settled")
+	}
+}
+
 func newTestGateway(t *testing.T, providers []config.ProviderConfig, routes []config.RouteConfig) (*httptest.Server, *captureUsageSink) {
+	return newTestGatewayWithBilling(t, providers, routes, nil)
+}
+
+func newTestGatewayWithBilling(t *testing.T, providers []config.ProviderConfig, routes []config.RouteConfig, meter billing.Meter) (*httptest.Server, *captureUsageSink) {
 	t.Helper()
 	authenticator, err := auth.NewStatic(`[{"key":"client-secret","key_id":"key-1","tenant_id":"tenant-1","project_id":"project-1","scopes":["inference"]}]`, false)
 	if err != nil {
@@ -201,6 +283,7 @@ func newTestGateway(t *testing.T, providers []config.ProviderConfig, routes []co
 		Router:        routing.New(modelCatalog),
 		Forwarder:     provider.New(cfg.UpstreamHTTP, metrics),
 		UsageSink:     sink,
+		BillingMeter:  meter,
 		Metrics:       metrics,
 		Logger:        slog.New(slog.NewTextHandler(io.Discard, nil)),
 		MaxBodyBytes:  1 << 20,