summaryrefslogtreecommitdiff
path: root/internal/httpapi
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--internal/httpapi/api.go113
-rw-r--r--internal/httpapi/api_test.go83
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,
generated by cgit v1.2.3 (git 2.46.0) at 2026-08-04 19:44:00 +0000