AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。 含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,139 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/apikey"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
var ErrAdmissionUnavailable = errors.New("admission control unavailable")
|
||||
|
||||
type AdmissionReason string
|
||||
|
||||
const (
|
||||
AdmissionAllowed AdmissionReason = ""
|
||||
AdmissionRateLimit AdmissionReason = "rate_limit"
|
||||
AdmissionMonthlyQuota AdmissionReason = "monthly_quota"
|
||||
)
|
||||
|
||||
type AdmissionDecision struct {
|
||||
Allowed bool
|
||||
Reason AdmissionReason
|
||||
Limit int64
|
||||
Remaining int64
|
||||
ResetAt time.Time
|
||||
RetryAfter time.Duration
|
||||
}
|
||||
|
||||
type AdmissionController interface {
|
||||
Allow(context.Context, apikey.Principal, time.Time) (AdmissionDecision, error)
|
||||
}
|
||||
|
||||
type RedisAdmissionController struct {
|
||||
client *redis.Client
|
||||
script *redis.Script
|
||||
}
|
||||
|
||||
func NewRedisAdmissionController(client *redis.Client) *RedisAdmissionController {
|
||||
return &RedisAdmissionController{client: client, script: redis.NewScript(admissionScript)}
|
||||
}
|
||||
|
||||
func (c *RedisAdmissionController) Allow(ctx context.Context, principal apikey.Principal, now time.Time) (AdmissionDecision, error) {
|
||||
if principal.RequestsPerMinute == 0 && principal.MonthlyRequestQuota == 0 {
|
||||
return AdmissionDecision{Allowed: true}, nil
|
||||
}
|
||||
if c == nil || c.client == nil || principal.APIKeyID == "" {
|
||||
return AdmissionDecision{}, ErrAdmissionUnavailable
|
||||
}
|
||||
now = now.UTC()
|
||||
minuteReset := now.Truncate(time.Minute).Add(time.Minute)
|
||||
monthReset := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, time.UTC)
|
||||
minuteKey := fmt.Sprintf("gateway:limit:api-key:%s:minute:%d", principal.APIKeyID, now.Unix()/60)
|
||||
monthKey := fmt.Sprintf("gateway:limit:api-key:%s:month:%s", principal.APIKeyID, now.Format("200601"))
|
||||
result, err := c.script.Run(ctx, c.client, []string{minuteKey, monthKey},
|
||||
principal.RequestsPerMinute, principal.MonthlyRequestQuota,
|
||||
int64(minuteReset.Sub(now).Seconds())+2, int64(monthReset.Sub(now).Seconds())+86400,
|
||||
).Slice()
|
||||
if err != nil {
|
||||
return AdmissionDecision{}, fmt.Errorf("%w: %v", ErrAdmissionUnavailable, err)
|
||||
}
|
||||
if len(result) != 3 {
|
||||
return AdmissionDecision{}, ErrAdmissionUnavailable
|
||||
}
|
||||
code, err := redisInteger(result[0])
|
||||
if err != nil {
|
||||
return AdmissionDecision{}, ErrAdmissionUnavailable
|
||||
}
|
||||
minuteCount, err := redisInteger(result[1])
|
||||
if err != nil {
|
||||
return AdmissionDecision{}, ErrAdmissionUnavailable
|
||||
}
|
||||
monthCount, err := redisInteger(result[2])
|
||||
if err != nil {
|
||||
return AdmissionDecision{}, ErrAdmissionUnavailable
|
||||
}
|
||||
switch code {
|
||||
case 0:
|
||||
decision := AdmissionDecision{Allowed: true}
|
||||
if principal.RequestsPerMinute > 0 {
|
||||
decision.Limit = int64(principal.RequestsPerMinute)
|
||||
decision.Remaining = max(decision.Limit-minuteCount, 0)
|
||||
decision.ResetAt = minuteReset
|
||||
}
|
||||
return decision, nil
|
||||
case 1:
|
||||
return AdmissionDecision{
|
||||
Allowed: false, Reason: AdmissionRateLimit, Limit: int64(principal.RequestsPerMinute),
|
||||
Remaining: 0, ResetAt: minuteReset, RetryAfter: minuteReset.Sub(now),
|
||||
}, nil
|
||||
case 2:
|
||||
return AdmissionDecision{
|
||||
Allowed: false, Reason: AdmissionMonthlyQuota, Limit: principal.MonthlyRequestQuota,
|
||||
Remaining: max(principal.MonthlyRequestQuota-monthCount, 0), ResetAt: monthReset, RetryAfter: monthReset.Sub(now),
|
||||
}, nil
|
||||
default:
|
||||
return AdmissionDecision{}, ErrAdmissionUnavailable
|
||||
}
|
||||
}
|
||||
|
||||
func redisInteger(value any) (int64, error) {
|
||||
switch number := value.(type) {
|
||||
case int64:
|
||||
return number, nil
|
||||
case string:
|
||||
return strconv.ParseInt(number, 10, 64)
|
||||
case []byte:
|
||||
return strconv.ParseInt(string(number), 10, 64)
|
||||
default:
|
||||
return 0, fmt.Errorf("unexpected redis integer %T", value)
|
||||
}
|
||||
}
|
||||
|
||||
const admissionScript = `
|
||||
local rpm = tonumber(ARGV[1])
|
||||
local monthly_quota = tonumber(ARGV[2])
|
||||
local minute_count = tonumber(redis.call('GET', KEYS[1]) or '0')
|
||||
local month_count = tonumber(redis.call('GET', KEYS[2]) or '0')
|
||||
|
||||
if monthly_quota > 0 and month_count >= monthly_quota then
|
||||
return {2, minute_count, month_count}
|
||||
end
|
||||
if rpm > 0 and minute_count >= rpm then
|
||||
return {1, minute_count, month_count}
|
||||
end
|
||||
|
||||
if rpm > 0 then
|
||||
minute_count = redis.call('INCR', KEYS[1])
|
||||
if minute_count == 1 then redis.call('EXPIRE', KEYS[1], tonumber(ARGV[3])) end
|
||||
end
|
||||
if monthly_quota > 0 then
|
||||
month_count = redis.call('INCR', KEYS[2])
|
||||
if month_count == 1 then redis.call('EXPIRE', KEYS[2], tonumber(ARGV[4])) end
|
||||
end
|
||||
return {0, minute_count, month_count}
|
||||
`
|
||||
@@ -0,0 +1,30 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestRedisInteger(t *testing.T) {
|
||||
for _, value := range []any{int64(7), "7", []byte("7")} {
|
||||
got, err := redisInteger(value)
|
||||
if err != nil || got != 7 {
|
||||
t.Fatalf("redisInteger(%T) = %d, %v", value, got, err)
|
||||
}
|
||||
}
|
||||
if _, err := redisInteger(7.0); err == nil {
|
||||
t.Fatal("unexpected redis number type accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdmissionResetBoundaries(t *testing.T) {
|
||||
now := time.Date(2026, time.August, 10, 13, 25, 40, 0, time.UTC)
|
||||
minuteReset := now.Truncate(time.Minute).Add(time.Minute)
|
||||
monthReset := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, time.UTC)
|
||||
if minuteReset != time.Date(2026, time.August, 10, 13, 26, 0, 0, time.UTC) {
|
||||
t.Fatal("minute boundary is incorrect")
|
||||
}
|
||||
if monthReset != time.Date(2026, time.September, 1, 0, 0, 0, 0, time.UTC) {
|
||||
t.Fatal("month boundary is incorrect")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,221 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/apikey"
|
||||
auditpkg "aigateway.local/core/internal/audit"
|
||||
"aigateway.local/core/internal/contentpolicy"
|
||||
"aigateway.local/core/internal/pricing"
|
||||
)
|
||||
|
||||
const auditRequestCaptureBytes = 64 << 10
|
||||
|
||||
type AuditRecorder interface {
|
||||
Record(auditpkg.Event) bool
|
||||
}
|
||||
|
||||
type auditSpan struct {
|
||||
recorder AuditRecorder
|
||||
started time.Time
|
||||
event auditpkg.Event
|
||||
mu sync.Mutex
|
||||
capture *captureReadCloser
|
||||
target string
|
||||
pricing *pricing.Service
|
||||
policies []contentpolicy.Match
|
||||
redacted bool
|
||||
}
|
||||
|
||||
func newAuditSpan(recorder AuditRecorder, principal apikey.Principal, request *http.Request, started time.Time) *auditSpan {
|
||||
if recorder == nil {
|
||||
return nil
|
||||
}
|
||||
event := auditpkg.Event{TenantID: principal.TenantID, RequestID: RequestID(request.Context()), Protocol: request.URL.Path, RecordedAt: time.Now().UTC()}
|
||||
if principal.APIKeyID != "" {
|
||||
id := principal.APIKeyID
|
||||
event.APIKeyID = &id
|
||||
}
|
||||
return &auditSpan{recorder: recorder, started: started, event: event}
|
||||
}
|
||||
|
||||
func (s *auditSpan) setModel(model string) {
|
||||
if s != nil && model != "" {
|
||||
s.mu.Lock()
|
||||
s.event.Model = model
|
||||
s.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *auditSpan) setRoute(providerCode, targetModel string) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.event.ProviderCode = providerCode
|
||||
s.target = targetModel
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *auditSpan) setUsage(usage TokenUsage) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.event.PromptTokens = usage.Input
|
||||
s.event.CompletionTokens = usage.Output
|
||||
s.calculateCostLocked()
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *auditSpan) setContentPolicy(result contentpolicy.Result) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.policies = append(s.policies, result.Matches...)
|
||||
s.redacted = result.Redacted
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *auditSpan) captureBody(body io.ReadCloser) io.ReadCloser {
|
||||
if s == nil || body == nil || body == http.NoBody {
|
||||
return body
|
||||
}
|
||||
capture := &captureReadCloser{ReadCloser: body, limit: auditRequestCaptureBytes}
|
||||
s.capture = capture
|
||||
return capture
|
||||
}
|
||||
|
||||
func (s *auditSpan) finish(status int) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
if s.event.Model == "" && s.capture != nil {
|
||||
var payload map[string]json.RawMessage
|
||||
if json.Unmarshal(s.capture.buffer, &payload) == nil {
|
||||
_ = json.Unmarshal(payload["model"], &s.event.Model)
|
||||
}
|
||||
}
|
||||
s.calculateCostLocked()
|
||||
s.event.StatusCode = status
|
||||
s.event.LatencyMS = int(time.Since(s.started).Milliseconds())
|
||||
s.event.RecordedAt = time.Now().UTC()
|
||||
labels := map[string]any{}
|
||||
if s.target != "" && s.target != s.event.Model {
|
||||
labels["target_model"] = s.target
|
||||
labels["routed"] = true
|
||||
}
|
||||
if len(s.policies) > 0 {
|
||||
labels["content_policies"] = s.policies
|
||||
}
|
||||
if s.redacted {
|
||||
labels["content_redacted"] = true
|
||||
}
|
||||
if s.event.PriceID != "" {
|
||||
labels["price_id"] = s.event.PriceID
|
||||
labels["currency"] = s.event.Currency
|
||||
}
|
||||
s.event.Labels = labels
|
||||
event := s.event
|
||||
s.mu.Unlock()
|
||||
s.recorder.Record(event)
|
||||
}
|
||||
|
||||
func (s *auditSpan) calculateCostLocked() {
|
||||
if s.pricing == nil || s.event.ProviderCode == "" || s.event.Model == "" {
|
||||
return
|
||||
}
|
||||
model := s.event.Model
|
||||
if s.target != "" {
|
||||
model = s.target
|
||||
}
|
||||
cost := s.pricing.Calculate(s.event.ProviderCode, model, s.event.PromptTokens, s.event.CompletionTokens, s.started.UTC())
|
||||
s.event.CostMicrounits = cost.Microunits
|
||||
s.event.PriceID = cost.PriceID
|
||||
s.event.Currency = cost.Currency
|
||||
}
|
||||
|
||||
type captureReadCloser struct {
|
||||
io.ReadCloser
|
||||
buffer []byte
|
||||
limit int
|
||||
}
|
||||
|
||||
func (r *captureReadCloser) Read(buffer []byte) (int, error) {
|
||||
n, err := r.ReadCloser.Read(buffer)
|
||||
if n > 0 && len(r.buffer) < r.limit {
|
||||
remaining := r.limit - len(r.buffer)
|
||||
r.buffer = append(r.buffer, buffer[:min(n, remaining)]...)
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
type statusResponseWriter struct {
|
||||
http.ResponseWriter
|
||||
status int
|
||||
}
|
||||
|
||||
func (w *statusResponseWriter) WriteHeader(status int) {
|
||||
if w.status != 0 {
|
||||
return
|
||||
}
|
||||
w.status = status
|
||||
w.ResponseWriter.WriteHeader(status)
|
||||
}
|
||||
|
||||
func (w *statusResponseWriter) Write(buffer []byte) (int, error) {
|
||||
if w.status == 0 {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}
|
||||
return w.ResponseWriter.Write(buffer)
|
||||
}
|
||||
|
||||
func (w *statusResponseWriter) Status() int {
|
||||
if w.status == 0 {
|
||||
return http.StatusOK
|
||||
}
|
||||
return w.status
|
||||
}
|
||||
|
||||
func (w *statusResponseWriter) Flush() {
|
||||
if w.status == 0 {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}
|
||||
if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
func (w *statusResponseWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
|
||||
|
||||
func (w *statusResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
||||
if hijacker, ok := w.ResponseWriter.(http.Hijacker); ok {
|
||||
return hijacker.Hijack()
|
||||
}
|
||||
return nil, nil, http.ErrNotSupported
|
||||
}
|
||||
|
||||
func (w *statusResponseWriter) Push(target string, options *http.PushOptions) error {
|
||||
if pusher, ok := w.ResponseWriter.(http.Pusher); ok {
|
||||
return pusher.Push(target, options)
|
||||
}
|
||||
return http.ErrNotSupported
|
||||
}
|
||||
|
||||
func (w *statusResponseWriter) ReadFrom(reader io.Reader) (int64, error) {
|
||||
if w.status == 0 {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}
|
||||
if readerFrom, ok := w.ResponseWriter.(io.ReaderFrom); ok {
|
||||
return readerFrom.ReadFrom(reader)
|
||||
}
|
||||
return io.Copy(struct{ io.Writer }{w.ResponseWriter}, reader)
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/audit"
|
||||
"aigateway.local/core/internal/contentpolicy"
|
||||
"aigateway.local/core/internal/pricing"
|
||||
provideropenai "aigateway.local/core/internal/provider/openai"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
func TestContentPolicyAndPricingIntegration(t *testing.T) {
|
||||
databaseURL := os.Getenv("CONTENT_PRICING_TEST_DATABASE_URL")
|
||||
if databaseURL == "" {
|
||||
t.Skip("CONTENT_PRICING_TEST_DATABASE_URL is not set")
|
||||
}
|
||||
ctx := context.Background()
|
||||
pool, err := pgxpool.New(ctx, databaseURL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer pool.Close()
|
||||
const blockID = "00000000-0000-4000-8000-000000001601"
|
||||
const priceID = "00000000-0000-4000-8000-000000001602"
|
||||
_, err = pool.Exec(ctx, `INSERT INTO gateway.content_policies(id,name,action,priority,rules,enabled) VALUES($1,'integration block','block',2000,'[{"name":"forbidden","pattern":"forbidden","replacement":"[BLOCKED]"}]',true) ON CONFLICT(id) DO UPDATE SET enabled=true`, blockID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = pool.Exec(ctx, `INSERT INTO gateway.model_prices(id,provider_code,model_pattern,input_microunits_per_million,output_microunits_per_million,currency,effective_from,enabled) VALUES($1,'environment','gpt-test',2000000,8000000,'USD',clock_timestamp()-interval '1 hour',true) ON CONFLICT(id) DO UPDATE SET enabled=true`, priceID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() {
|
||||
_, _ = pool.Exec(context.Background(), `DELETE FROM gateway.outbox_events WHERE event_type='content_policy.matched' AND aggregate_id IN ('m3-content-price-redact','m3-content-price-block')`)
|
||||
_, _ = pool.Exec(context.Background(), `DELETE FROM gateway.content_policies WHERE id=$1`, blockID)
|
||||
_, _ = pool.Exec(context.Background(), `DELETE FROM gateway.model_prices WHERE id=$1`, priceID)
|
||||
}()
|
||||
|
||||
var upstreamCalls atomic.Int64
|
||||
bodySeen := make(chan string, 1)
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCalls.Add(1)
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
bodySeen <- string(body)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `{"id":"ok","usage":{"prompt_tokens":1000,"completion_tokens":500,"total_tokens":1500}}`)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
adapter, _ := provideropenai.New(upstream.URL, "upstream-key")
|
||||
engine := contentpolicy.NewEngine(pool, time.Minute, slog.Default())
|
||||
if err := engine.Reload(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
prices := pricing.NewService(pool, time.Minute, slog.Default())
|
||||
if err := prices.Reload(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
recorder := audit.NewRecorder(pool, slog.Default(), 100, 1, 10*time.Millisecond)
|
||||
recordCtx, cancel := context.WithCancel(ctx)
|
||||
stopped := make(chan struct{})
|
||||
go func() { recorder.Run(recordCtx); close(stopped) }()
|
||||
defer func() { cancel(); <-stopped }()
|
||||
proxy := NewProxy(adapter, "test-key", 1<<20, slog.Default())
|
||||
proxy.SetAuditRecorder(recorder)
|
||||
proxy.SetContentPolicyEngine(engine)
|
||||
proxy.SetPricingService(prices)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
proxy.ServeHTTP(w, r.WithContext(WithRequestID(r.Context(), r.Header.Get("X-Request-ID"))))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
req, _ := http.NewRequest(http.MethodPost, server.URL+"/v1/chat/completions", strings.NewReader(`{"model":"gpt-test","messages":[{"role":"user","content":"token sk-1234567890123456"}]}`))
|
||||
req.Header.Set("Authorization", "Bearer test-key")
|
||||
req.Header.Set("X-Request-ID", "m3-content-price-redact")
|
||||
response, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, _ = io.ReadAll(response.Body)
|
||||
_ = response.Body.Close()
|
||||
if response.StatusCode != 200 || response.Header.Get("X-Gateway-Content-Redacted") != "true" {
|
||||
t.Fatalf("unexpected response %d headers=%v", response.StatusCode, response.Header)
|
||||
}
|
||||
select {
|
||||
case body := <-bodySeen:
|
||||
if strings.Contains(body, "sk-1234567890123456") || !strings.Contains(body, "[REDACTED_API_KEY]") {
|
||||
t.Fatalf("upstream saw unsafe body %s", body)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("upstream did not receive request")
|
||||
}
|
||||
|
||||
blocked, _ := http.NewRequest(http.MethodPost, server.URL+"/v1/responses", strings.NewReader(`{"model":"gpt-test","input":"forbidden"}`))
|
||||
blocked.Header.Set("Authorization", "Bearer test-key")
|
||||
blocked.Header.Set("X-Request-ID", "m3-content-price-block")
|
||||
blockedResponse, err := http.DefaultClient.Do(blocked)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, _ = io.ReadAll(blockedResponse.Body)
|
||||
_ = blockedResponse.Body.Close()
|
||||
if blockedResponse.StatusCode != http.StatusUnprocessableEntity || upstreamCalls.Load() != 1 {
|
||||
t.Fatalf("block failed status=%d calls=%d", blockedResponse.StatusCode, upstreamCalls.Load())
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for {
|
||||
var cost *int64
|
||||
var labels string
|
||||
err = pool.QueryRow(ctx, `SELECT cost_microunits,labels::text FROM gateway.audit_events WHERE request_id='m3-content-price-redact' ORDER BY recorded_at DESC LIMIT 1`).Scan(&cost, &labels)
|
||||
if err == nil && cost != nil && *cost == 6000 && strings.Contains(labels, "content_redacted") {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("audit cost/redaction missing cost=%v labels=%s err=%v", cost, labels, err)
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
deadline = time.Now().Add(2 * time.Second)
|
||||
for {
|
||||
var count int
|
||||
err = pool.QueryRow(ctx, `SELECT count(*) FROM gateway.outbox_events WHERE event_type='content_policy.matched' AND aggregate_id IN ('m3-content-price-redact','m3-content-price-block')`).Scan(&count)
|
||||
if err == nil && count == 2 {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("content policy notification outbox missing count=%d err=%v", count, err)
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
package gateway
|
||||
|
||||
import "context"
|
||||
|
||||
type requestIDKey struct{}
|
||||
|
||||
func WithRequestID(ctx context.Context, requestID string) context.Context {
|
||||
return context.WithValue(ctx, requestIDKey{}, requestID)
|
||||
}
|
||||
|
||||
func RequestID(ctx context.Context) string {
|
||||
value, _ := ctx.Value(requestIDKey{}).(string)
|
||||
return value
|
||||
}
|
||||
@@ -0,0 +1,430 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/subtle"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/apikey"
|
||||
"aigateway.local/core/internal/contentpolicy"
|
||||
"aigateway.local/core/internal/pricing"
|
||||
"aigateway.local/core/internal/provider"
|
||||
)
|
||||
|
||||
type Proxy struct {
|
||||
resolver AdapterResolver
|
||||
auth apikey.KeyAuthenticator
|
||||
maxBody int64
|
||||
logger *slog.Logger
|
||||
transport *http.Transport
|
||||
proxies sync.Map
|
||||
circuits sync.Map
|
||||
admission AdmissionController
|
||||
tokenQuota TokenQuotaController
|
||||
resilience ResiliencePolicy
|
||||
audit AuditRecorder
|
||||
policies *contentpolicy.Engine
|
||||
pricing *pricing.Service
|
||||
}
|
||||
|
||||
type cachedProxy struct {
|
||||
key string
|
||||
proxy *httputil.ReverseProxy
|
||||
}
|
||||
|
||||
var (
|
||||
ErrProviderNotFound = errors.New("requested provider is not available")
|
||||
ErrProviderUnavailable = errors.New("provider configuration is unavailable")
|
||||
)
|
||||
|
||||
type ResolvedAdapter struct {
|
||||
Code string
|
||||
Revision int64
|
||||
Adapter provider.Adapter
|
||||
Capabilities map[provider.Capability]bool
|
||||
}
|
||||
|
||||
type AdapterResolver interface {
|
||||
Resolve(providerCode string) (ResolvedAdapter, error)
|
||||
}
|
||||
|
||||
func NewProxy(adapter provider.Adapter, apiKey string, maxBody int64, logger *slog.Logger) *Proxy {
|
||||
return NewProxyWithAuthenticator(adapter, staticKeyAuthenticator(apiKey), maxBody, logger)
|
||||
}
|
||||
|
||||
func NewProxyWithAuthenticator(adapter provider.Adapter, authenticator apikey.KeyAuthenticator, maxBody int64, logger *slog.Logger) *Proxy {
|
||||
return NewDynamicProxy(staticAdapterResolver{adapter: adapter}, authenticator, maxBody, logger)
|
||||
}
|
||||
|
||||
func NewDynamicProxy(resolver AdapterResolver, authenticator apikey.KeyAuthenticator, maxBody int64, logger *slog.Logger) *Proxy {
|
||||
transport := &http.Transport{
|
||||
Proxy: http.ProxyFromEnvironment,
|
||||
DialContext: (&net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
|
||||
ForceAttemptHTTP2: true,
|
||||
MaxIdleConns: 512,
|
||||
MaxIdleConnsPerHost: 256,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
TLSHandshakeTimeout: 5 * time.Second,
|
||||
ResponseHeaderTimeout: 60 * time.Second,
|
||||
ExpectContinueTimeout: time.Second,
|
||||
}
|
||||
return &Proxy{resolver: resolver, auth: authenticator, maxBody: maxBody, logger: logger, transport: transport, resilience: DefaultResiliencePolicy()}
|
||||
}
|
||||
|
||||
func (p *Proxy) SetAdmissionController(controller AdmissionController) {
|
||||
p.admission = controller
|
||||
}
|
||||
|
||||
func (p *Proxy) SetTokenQuotaController(controller TokenQuotaController) {
|
||||
p.tokenQuota = controller
|
||||
}
|
||||
|
||||
func (p *Proxy) SetResiliencePolicy(policy ResiliencePolicy) {
|
||||
p.resilience = policy
|
||||
p.transport.ResponseHeaderTimeout = policy.ResponseHeaderTimeout
|
||||
}
|
||||
|
||||
func (p *Proxy) SetAuditRecorder(recorder AuditRecorder) { p.audit = recorder }
|
||||
func (p *Proxy) SetContentPolicyEngine(engine *contentpolicy.Engine) { p.policies = engine }
|
||||
func (p *Proxy) SetPricingService(service *pricing.Service) { p.pricing = service }
|
||||
|
||||
func (p *Proxy) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
|
||||
started := time.Now()
|
||||
if !isSupportedPath(request.URL.Path, request.Method) {
|
||||
writeOpenAIError(writer, http.StatusNotFound, "invalid_request_error", "unsupported gateway endpoint")
|
||||
return
|
||||
}
|
||||
principal, err := p.authorized(request)
|
||||
if err != nil {
|
||||
if errors.Is(err, apikey.ErrStore) {
|
||||
writeOpenAIError(writer, http.StatusServiceUnavailable, "authentication_unavailable", "API key service is unavailable")
|
||||
return
|
||||
}
|
||||
writer.Header().Set("WWW-Authenticate", "Bearer")
|
||||
writeOpenAIError(writer, http.StatusUnauthorized, "authentication_error", "invalid API key")
|
||||
return
|
||||
}
|
||||
span := newAuditSpan(p.audit, principal, request, started)
|
||||
if span != nil {
|
||||
span.pricing = p.pricing
|
||||
}
|
||||
if span != nil {
|
||||
statusWriter := &statusResponseWriter{ResponseWriter: writer}
|
||||
writer = statusWriter
|
||||
defer func() { span.finish(statusWriter.Status()) }()
|
||||
}
|
||||
if p.admission != nil && principal.APIKeyID != "" {
|
||||
decision, err := p.admission.Allow(request.Context(), principal, time.Now())
|
||||
if err != nil {
|
||||
writeOpenAIError(writer, http.StatusServiceUnavailable, "rate_limit_unavailable", "rate limit service is unavailable")
|
||||
return
|
||||
}
|
||||
writeAdmissionHeaders(writer.Header(), decision)
|
||||
if !decision.Allowed {
|
||||
writer.Header().Set("Retry-After", strconv.FormatInt(max(int64(decision.RetryAfter.Seconds()), 1), 10))
|
||||
if decision.Reason == AdmissionMonthlyQuota {
|
||||
writeOpenAIError(writer, http.StatusTooManyRequests, "insufficient_quota", "monthly request quota exceeded")
|
||||
} else {
|
||||
writeOpenAIError(writer, http.StatusTooManyRequests, "rate_limit_error", "request rate limit exceeded")
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
if request.ContentLength > p.maxBody {
|
||||
writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
|
||||
return
|
||||
}
|
||||
if p.policies != nil {
|
||||
policyResult, policyErr := p.policies.Apply(request, p.maxBody, principal.APIKeyID)
|
||||
if errors.Is(policyErr, contentpolicy.ErrBodyTooLarge) {
|
||||
writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
|
||||
return
|
||||
}
|
||||
if policyErr != nil {
|
||||
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request body could not be inspected")
|
||||
return
|
||||
}
|
||||
span.setContentPolicy(policyResult)
|
||||
if policyResult.Blocked {
|
||||
writeOpenAIError(writer, http.StatusUnprocessableEntity, "content_policy_violation", "request was blocked by content policy")
|
||||
return
|
||||
}
|
||||
if policyResult.Redacted {
|
||||
writer.Header().Set("X-Gateway-Content-Redacted", "true")
|
||||
}
|
||||
}
|
||||
routeResolver, routingEnabled := p.resolver.(ModelRouteResolver)
|
||||
routingEnabled = routingEnabled && routeResolver.ModelRoutingEnabled()
|
||||
var modelPayload modelRequest
|
||||
if routingEnabled {
|
||||
modelPayload, err = readModelRequest(request, p.maxBody)
|
||||
if errors.Is(err, errRequestBodyTooLarge) {
|
||||
writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request body could not be read")
|
||||
return
|
||||
}
|
||||
}
|
||||
var usage *usageSession
|
||||
if p.tokenQuota != nil && principal.APIKeyID != "" {
|
||||
estimate := int64(0)
|
||||
if principal.MonthlyTokenQuota > 0 {
|
||||
estimate, err = prepareTokenBudget(request, p.maxBody)
|
||||
if errors.Is(err, errRequestBodyTooLarge) {
|
||||
writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request body could not be read")
|
||||
return
|
||||
}
|
||||
}
|
||||
reservation, reserveErr := p.tokenQuota.Reserve(request.Context(), principal, estimate, time.Now())
|
||||
if reserveErr != nil && principal.MonthlyTokenQuota > 0 {
|
||||
writeOpenAIError(writer, http.StatusServiceUnavailable, "token_quota_unavailable", "token quota service is unavailable")
|
||||
return
|
||||
}
|
||||
if reserveErr == nil && !reservation.Allowed {
|
||||
writeTokenQuotaHeaders(writer.Header(), reservation)
|
||||
writer.Header().Set("Retry-After", strconv.FormatInt(max(int64(time.Until(reservation.ResetAt).Seconds()), 1), 10))
|
||||
writeOpenAIError(writer, http.StatusTooManyRequests, "insufficient_quota", "monthly token quota exceeded")
|
||||
return
|
||||
}
|
||||
if reserveErr == nil && reservation.CounterKey != "" {
|
||||
writeTokenQuotaHeaders(writer.Header(), reservation)
|
||||
usage = &usageSession{controller: p.tokenQuota, reservation: reservation, log: p.logger}
|
||||
request = withUsageSession(request, usage)
|
||||
defer usage.finish(TokenUsage{})
|
||||
}
|
||||
}
|
||||
if span != nil {
|
||||
if usage == nil {
|
||||
usage = &usageSession{log: p.logger}
|
||||
request = withUsageSession(request, usage)
|
||||
defer usage.finish(TokenUsage{})
|
||||
}
|
||||
usage.onFinish = span.setUsage
|
||||
}
|
||||
if request.Body != nil && request.Method != http.MethodGet {
|
||||
if request.Header.Get("Idempotency-Key") != "" && request.GetBody == nil {
|
||||
if _, err := prepareTokenBudget(request, p.maxBody); err != nil {
|
||||
if errors.Is(err, errRequestBodyTooLarge) {
|
||||
writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
|
||||
} else {
|
||||
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request body could not be read")
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
request.Body = http.MaxBytesReader(writer, request.Body, p.maxBody)
|
||||
}
|
||||
providerCode := strings.ToLower(strings.TrimSpace(request.Header.Get("X-Gateway-Provider")))
|
||||
var resolved ResolvedAdapter
|
||||
if routingEnabled && modelPayload.model != "" {
|
||||
span.setModel(modelPayload.model)
|
||||
tenantID := ""
|
||||
if principal.TenantID != nil {
|
||||
tenantID = *principal.TenantID
|
||||
}
|
||||
route, routeErr := routeResolver.ResolveModelRoute(ModelRouteQuery{
|
||||
ProviderCode: providerCode, Model: modelPayload.model, Endpoint: request.URL.Path,
|
||||
APIKeyID: principal.APIKeyID, TenantID: tenantID, Seed: RequestID(request.Context()),
|
||||
})
|
||||
if routeErr != nil {
|
||||
writeOpenAIError(writer, http.StatusServiceUnavailable, "provider_unavailable", "model routing configuration is unavailable")
|
||||
return
|
||||
}
|
||||
if route.Known && !route.Matched {
|
||||
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "model route is not available for this request")
|
||||
return
|
||||
}
|
||||
if route.Matched {
|
||||
resolved = route.ResolvedAdapter
|
||||
if err := modelPayload.rewrite(request, route.TargetModel); err != nil {
|
||||
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request model could not be rewritten")
|
||||
return
|
||||
}
|
||||
writer.Header().Set("X-Gateway-Model", route.TargetModel)
|
||||
}
|
||||
}
|
||||
if resolved.Adapter == nil {
|
||||
resolved, err = p.resolver.Resolve(providerCode)
|
||||
}
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrProviderNotFound) {
|
||||
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "requested provider is not available")
|
||||
return
|
||||
}
|
||||
writeOpenAIError(writer, http.StatusServiceUnavailable, "provider_unavailable", "provider configuration is unavailable")
|
||||
return
|
||||
}
|
||||
capability := capabilityForPath(request.URL.Path)
|
||||
if capability != "" && !resolved.Capabilities[capability] {
|
||||
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "provider does not support this endpoint")
|
||||
return
|
||||
}
|
||||
request.Header.Del("X-Gateway-Provider")
|
||||
writer.Header().Set("X-Gateway-Provider", resolved.Code)
|
||||
span.setRoute(resolved.Code, writer.Header().Get("X-Gateway-Model"))
|
||||
request.Body = span.captureBody(request.Body)
|
||||
p.proxyFor(resolved).ServeHTTP(writer, request)
|
||||
}
|
||||
|
||||
func (p *Proxy) proxyFor(resolved ResolvedAdapter) *httputil.ReverseProxy {
|
||||
target := resolved.Adapter.Target()
|
||||
key := fmt.Sprintf("%s:%d:%s", resolved.Code, resolved.Revision, target.String())
|
||||
if cached, ok := p.proxies.Load(resolved.Code); ok {
|
||||
entry := cached.(cachedProxy)
|
||||
if entry.key == key {
|
||||
return entry.proxy
|
||||
}
|
||||
}
|
||||
reverseProxy := httputil.NewSingleHostReverseProxy(target)
|
||||
originalDirector := reverseProxy.Director
|
||||
reverseProxy.Director = func(request *http.Request) {
|
||||
originalDirector(request)
|
||||
request.Host = target.Host
|
||||
resolved.Adapter.Prepare(request)
|
||||
}
|
||||
reverseProxy.FlushInterval = -1
|
||||
circuitValue, _ := p.circuits.LoadOrStore(resolved.Code, newCircuitBreaker(p.resilience))
|
||||
reverseProxy.Transport = &resilientTransport{
|
||||
base: p.transport, circuit: circuitValue.(*circuitBreaker), maxRetries: p.resilience.MaxRetries, backoff: p.resilience.RetryBackoff,
|
||||
}
|
||||
reverseProxy.ModifyResponse = func(response *http.Response) error {
|
||||
if session := usageSessionFrom(response.Request); session != nil {
|
||||
if response.StatusCode >= http.StatusOK && response.StatusCode < http.StatusMultipleChoices {
|
||||
session.fallback = session.reservation.Reserved
|
||||
response.Body = newUsageReadCloser(response.Body, response.Header.Get("Content-Type"), session)
|
||||
} else {
|
||||
session.fallback = 0
|
||||
session.finish(TokenUsage{})
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
reverseProxy.ErrorHandler = func(writer http.ResponseWriter, request *http.Request, err error) {
|
||||
if session := usageSessionFrom(request); session != nil {
|
||||
session.fallback = 0
|
||||
session.finish(TokenUsage{})
|
||||
}
|
||||
p.logger.Error("upstream request failed", "request_id", RequestID(request.Context()), "provider", resolved.Code, "error", err)
|
||||
if errors.Is(err, ErrCircuitOpen) {
|
||||
writer.Header().Set("Retry-After", "30")
|
||||
writeOpenAIError(writer, http.StatusServiceUnavailable, "provider_unavailable", "provider circuit is temporarily open")
|
||||
return
|
||||
}
|
||||
writeOpenAIError(writer, http.StatusBadGateway, "upstream_error", "upstream service is unavailable")
|
||||
}
|
||||
p.proxies.Store(resolved.Code, cachedProxy{key: key, proxy: reverseProxy})
|
||||
return reverseProxy
|
||||
}
|
||||
|
||||
func (p *Proxy) authorized(request *http.Request) (apikey.Principal, error) {
|
||||
presented := strings.TrimSpace(request.Header.Get("X-Gateway-API-Key"))
|
||||
if presented == "" {
|
||||
authorization := strings.TrimSpace(request.Header.Get("Authorization"))
|
||||
if len(authorization) > len("Bearer ") && strings.EqualFold(authorization[:len("Bearer ")], "Bearer ") {
|
||||
presented = strings.TrimSpace(authorization[len("Bearer "):])
|
||||
}
|
||||
}
|
||||
if p.auth == nil {
|
||||
return apikey.Principal{}, apikey.ErrInvalid
|
||||
}
|
||||
if authenticator, ok := p.auth.(apikey.PrincipalAuthenticator); ok {
|
||||
return authenticator.AuthenticatePrincipal(request.Context(), presented)
|
||||
}
|
||||
return apikey.Principal{}, p.auth.Authenticate(request.Context(), presented)
|
||||
}
|
||||
|
||||
func writeAdmissionHeaders(header http.Header, decision AdmissionDecision) {
|
||||
if decision.Limit <= 0 || decision.ResetAt.IsZero() {
|
||||
return
|
||||
}
|
||||
header.Set("X-RateLimit-Limit", strconv.FormatInt(decision.Limit, 10))
|
||||
header.Set("X-RateLimit-Remaining", strconv.FormatInt(max(decision.Remaining, 0), 10))
|
||||
header.Set("X-RateLimit-Reset", strconv.FormatInt(decision.ResetAt.Unix(), 10))
|
||||
}
|
||||
|
||||
type staticKeyAuthenticator string
|
||||
|
||||
func (a staticKeyAuthenticator) Authenticate(_ context.Context, presented string) error {
|
||||
key := string(a)
|
||||
if key == "" {
|
||||
return nil
|
||||
}
|
||||
if len(presented) != len(key) || subtle.ConstantTimeCompare([]byte(presented), []byte(key)) != 1 {
|
||||
return apikey.ErrInvalid
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type staticAdapterResolver struct{ adapter provider.Adapter }
|
||||
|
||||
func (r staticAdapterResolver) Resolve(code string) (ResolvedAdapter, error) {
|
||||
if r.adapter == nil || code != "" && code != "environment" {
|
||||
return ResolvedAdapter{}, ErrProviderNotFound
|
||||
}
|
||||
capabilities := make(map[provider.Capability]bool)
|
||||
for _, capability := range r.adapter.Capabilities() {
|
||||
capabilities[capability] = true
|
||||
}
|
||||
return ResolvedAdapter{Code: "environment", Adapter: r.adapter, Capabilities: capabilities}, nil
|
||||
}
|
||||
|
||||
func capabilityForPath(path string) provider.Capability {
|
||||
switch path {
|
||||
case "/v1/models":
|
||||
return provider.CapabilityModels
|
||||
case "/v1/chat/completions":
|
||||
return provider.CapabilityChat
|
||||
case "/v1/responses":
|
||||
return provider.CapabilityResponses
|
||||
case "/v1/embeddings":
|
||||
return provider.CapabilityEmbeddings
|
||||
case "/v1/messages":
|
||||
return provider.CapabilityMessages
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func isSupportedPath(path, method string) bool {
|
||||
switch path {
|
||||
case "/v1/models":
|
||||
return method == http.MethodGet
|
||||
case "/v1/chat/completions", "/v1/responses", "/v1/embeddings", "/v1/messages":
|
||||
return method == http.MethodPost
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func writeOpenAIError(writer http.ResponseWriter, status int, errorType, message string) {
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
writer.WriteHeader(status)
|
||||
_ = json.NewEncoder(writer).Encode(map[string]any{
|
||||
"error": map[string]any{"message": message, "type": errorType, "param": nil, "code": nil},
|
||||
})
|
||||
}
|
||||
|
||||
var errInvalidAdapter = errors.New("invalid provider adapter")
|
||||
|
||||
func ValidateAdapter(adapter provider.Adapter) error {
|
||||
if adapter == nil || adapter.Target() == nil || adapter.Name() == "" {
|
||||
return errInvalidAdapter
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/apikey"
|
||||
auditpkg "aigateway.local/core/internal/audit"
|
||||
"aigateway.local/core/internal/provider"
|
||||
provideropenai "aigateway.local/core/internal/provider/openai"
|
||||
)
|
||||
|
||||
type principalAuthenticator struct{ principal apikey.Principal }
|
||||
|
||||
func (a principalAuthenticator) Authenticate(context.Context, string) error { return nil }
|
||||
func (a principalAuthenticator) AuthenticatePrincipal(context.Context, string) (apikey.Principal, error) {
|
||||
return a.principal, nil
|
||||
}
|
||||
|
||||
type fixedAdmission struct{ decision AdmissionDecision }
|
||||
|
||||
func (a fixedAdmission) Allow(context.Context, apikey.Principal, time.Time) (AdmissionDecision, error) {
|
||||
return a.decision, nil
|
||||
}
|
||||
|
||||
type recordingTokenQuota struct {
|
||||
reservation TokenReservation
|
||||
estimate int64
|
||||
actual int64
|
||||
}
|
||||
|
||||
type fixedRoutingResolver struct {
|
||||
adapter ResolvedAdapter
|
||||
}
|
||||
|
||||
type recordingAudit struct{ event auditpkg.Event }
|
||||
|
||||
func (r *recordingAudit) Record(event auditpkg.Event) bool { r.event = event; return true }
|
||||
|
||||
func (r fixedRoutingResolver) Resolve(string) (ResolvedAdapter, error) { return r.adapter, nil }
|
||||
func (r fixedRoutingResolver) ModelRoutingEnabled() bool { return true }
|
||||
func (r fixedRoutingResolver) ResolveModelRoute(query ModelRouteQuery) (ModelRouteResult, error) {
|
||||
if query.Model == "public-chat" {
|
||||
return ModelRouteResult{ResolvedAdapter: r.adapter, TargetModel: "upstream-chat", Matched: true, Known: true}, nil
|
||||
}
|
||||
return ModelRouteResult{}, nil
|
||||
}
|
||||
|
||||
func (q *recordingTokenQuota) Reserve(_ context.Context, _ apikey.Principal, estimate int64, _ time.Time) (TokenReservation, error) {
|
||||
q.estimate = estimate
|
||||
return q.reservation, nil
|
||||
}
|
||||
|
||||
func (q *recordingTokenQuota) Commit(_ context.Context, _ TokenReservation, actual int64) error {
|
||||
q.actual = actual
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestProxyRejectsInvalidKey(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
|
||||
t.Fatal("upstream must not be called")
|
||||
}))
|
||||
defer upstream.Close()
|
||||
adapter, err := provideropenai.New(upstream.URL, "upstream-secret")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
proxy := NewProxy(adapter, "gateway-secret", 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
request.Header.Set("Authorization", "Bearer wrong")
|
||||
response := httptest.NewRecorder()
|
||||
proxy.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("got %d, want %d", response.Code, http.StatusUnauthorized)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyRejectsKnownOversizedBody(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
|
||||
t.Fatal("upstream must not be called")
|
||||
}))
|
||||
defer upstream.Close()
|
||||
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
|
||||
proxy := NewProxy(adapter, "gateway-secret", 4, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader("12345"))
|
||||
request.Header.Set("Authorization", "Bearer gateway-secret")
|
||||
response := httptest.NewRecorder()
|
||||
proxy.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusRequestEntityTooLarge {
|
||||
t.Fatalf("got %d, want %d", response.Code, http.StatusRequestEntityTooLarge)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyReplacesClientAuthorization(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
if got := request.Header.Get("Authorization"); got != "Bearer upstream-secret" {
|
||||
t.Fatalf("unexpected upstream authorization: %q", got)
|
||||
}
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = writer.Write([]byte(`{"ok":true}`))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
|
||||
proxy := NewProxy(adapter, "gateway-secret", 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
request.Header.Set("Authorization", "Bearer gateway-secret")
|
||||
response := httptest.NewRecorder()
|
||||
proxy.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusOK {
|
||||
t.Fatalf("got %d, want %d", response.Code, http.StatusOK)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyRejectsRateLimitedPrincipalBeforeUpstream(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
|
||||
t.Fatal("upstream must not be called")
|
||||
}))
|
||||
defer upstream.Close()
|
||||
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
|
||||
proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
|
||||
APIKeyID: "key-1", Scopes: []string{"gateway:invoke"}, RequestsPerMinute: 2,
|
||||
}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
proxy.SetAdmissionController(fixedAdmission{decision: AdmissionDecision{
|
||||
Allowed: false, Reason: AdmissionRateLimit, Limit: 2, Remaining: 0,
|
||||
ResetAt: time.Now().Add(time.Minute), RetryAfter: 30 * time.Second,
|
||||
}})
|
||||
|
||||
request := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
request.Header.Set("Authorization", "Bearer gateway-secret")
|
||||
response := httptest.NewRecorder()
|
||||
proxy.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("got %d, want %d", response.Code, http.StatusTooManyRequests)
|
||||
}
|
||||
if response.Header().Get("X-RateLimit-Limit") != "2" || response.Header().Get("Retry-After") == "" {
|
||||
t.Fatalf("missing rate limit headers: %#v", response.Header())
|
||||
}
|
||||
if !strings.Contains(response.Body.String(), `"type":"rate_limit_error"`) {
|
||||
t.Fatalf("unexpected body: %s", response.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyReconcilesReservedTokensWithUpstreamUsage(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = writer.Write([]byte(`{"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":5,"total_tokens":12}}`))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
|
||||
quota := &recordingTokenQuota{reservation: TokenReservation{
|
||||
Allowed: true, CounterKey: "tokens", Reserved: 24, Limit: 1000, Remaining: 976, ResetAt: time.Now().Add(time.Hour),
|
||||
}}
|
||||
proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
|
||||
APIKeyID: "key-1", Scopes: []string{"gateway:invoke"}, MonthlyTokenQuota: 1000,
|
||||
}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
proxy.SetTokenQuotaController(quota)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"test","max_tokens":20}`))
|
||||
request.Header.Set("Authorization", "Bearer gateway-secret")
|
||||
response := httptest.NewRecorder()
|
||||
proxy.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusOK {
|
||||
t.Fatalf("got %d, want %d: %s", response.Code, http.StatusOK, response.Body.String())
|
||||
}
|
||||
if quota.estimate <= 20 {
|
||||
t.Fatalf("request budget did not include input tokens: %d", quota.estimate)
|
||||
}
|
||||
if quota.actual != 12 {
|
||||
t.Fatalf("committed %d tokens, want 12", quota.actual)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyRejectsTokenReservationBeforeUpstream(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
|
||||
t.Fatal("upstream must not be called")
|
||||
}))
|
||||
defer upstream.Close()
|
||||
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
|
||||
quota := &recordingTokenQuota{reservation: TokenReservation{
|
||||
Allowed: false, Limit: 100, Remaining: 4, ResetAt: time.Now().Add(time.Hour),
|
||||
}}
|
||||
proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
|
||||
APIKeyID: "key-1", Scopes: []string{"gateway:invoke"}, MonthlyTokenQuota: 100,
|
||||
}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
proxy.SetTokenQuotaController(quota)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"test","max_output_tokens":20}`))
|
||||
request.Header.Set("Authorization", "Bearer gateway-secret")
|
||||
response := httptest.NewRecorder()
|
||||
proxy.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusTooManyRequests || !strings.Contains(response.Body.String(), `"type":"insufficient_quota"`) {
|
||||
t.Fatalf("unexpected response %d: %s", response.Code, response.Body.String())
|
||||
}
|
||||
if response.Header().Get("X-TokenLimit-Limit") != "100" {
|
||||
t.Fatalf("missing token quota headers: %#v", response.Header())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyRewritesModelAliasBeforeUpstream(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
body, _ := io.ReadAll(request.Body)
|
||||
if !strings.Contains(string(body), `"model":"upstream-chat"`) {
|
||||
t.Fatalf("model alias was not rewritten: %s", body)
|
||||
}
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = writer.Write([]byte(`{"choices":[]}`))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
|
||||
resolved := ResolvedAdapter{Code: "routed", Adapter: adapter, Capabilities: map[provider.Capability]bool{provider.CapabilityChat: true}}
|
||||
proxy := NewDynamicProxy(fixedRoutingResolver{adapter: resolved}, principalAuthenticator{principal: apikey.Principal{Scopes: []string{"gateway:invoke"}}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"public-chat","messages":[]}`))
|
||||
request.Header.Set("Authorization", "Bearer gateway-secret")
|
||||
response := httptest.NewRecorder()
|
||||
proxy.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusOK || response.Header().Get("X-Gateway-Model") != "upstream-chat" || response.Header().Get("X-Gateway-Provider") != "routed" {
|
||||
t.Fatalf("unexpected routed response %d %#v: %s", response.Code, response.Header(), response.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyRecordsAuditWithoutBufferingWholeResponse(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = writer.Write([]byte(`{"choices":[],"usage":{"prompt_tokens":9,"completion_tokens":4,"total_tokens":13}}`))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
|
||||
recorder := &recordingAudit{}
|
||||
proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
|
||||
APIKeyID: "11111111-1111-4111-8111-111111111111", Scopes: []string{"gateway:invoke"},
|
||||
}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
proxy.SetAuditRecorder(recorder)
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"audit-model","messages":[]}`))
|
||||
request.Header.Set("Authorization", "Bearer gateway-secret")
|
||||
response := httptest.NewRecorder()
|
||||
proxy.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusOK {
|
||||
t.Fatalf("unexpected response %d: %s", response.Code, response.Body.String())
|
||||
}
|
||||
if recorder.event.Model != "audit-model" || recorder.event.ProviderCode != "environment" || recorder.event.StatusCode != 200 || recorder.event.PromptTokens != 9 || recorder.event.CompletionTokens != 4 {
|
||||
t.Fatalf("unexpected audit event: %#v", recorder.event)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
var errRequestBodyTooLarge = errors.New("request body is too large")
|
||||
|
||||
func prepareTokenBudget(request *http.Request, maxBody int64) (int64, error) {
|
||||
if request.Body == nil || request.Method == http.MethodGet {
|
||||
return 0, nil
|
||||
}
|
||||
limited := io.LimitReader(request.Body, maxBody+1)
|
||||
body, err := io.ReadAll(limited)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
_ = request.Body.Close()
|
||||
if int64(len(body)) > maxBody {
|
||||
return 0, errRequestBodyTooLarge
|
||||
}
|
||||
restoreRequestBody(request, body)
|
||||
|
||||
inputEstimate := max(int64((len(body)+3)/4), 1)
|
||||
if request.URL.Path == "/v1/embeddings" {
|
||||
return inputEstimate, nil
|
||||
}
|
||||
outputBudget := int64(1024)
|
||||
decoder := json.NewDecoder(bytes.NewReader(body))
|
||||
decoder.UseNumber()
|
||||
var payload map[string]any
|
||||
if decoder.Decode(&payload) == nil {
|
||||
for _, field := range []string{"max_output_tokens", "max_completion_tokens", "max_tokens"} {
|
||||
if value := jsonInt64(payload[field]); value > 0 {
|
||||
outputBudget = value
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
// A per-request ceiling prevents malicious payloads from reserving unbounded Redis counters.
|
||||
return min(inputEstimate+outputBudget, int64(100_000_000)), nil
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
var ErrCircuitOpen = errors.New("provider circuit is open")
|
||||
|
||||
type ResiliencePolicy struct {
|
||||
ResponseHeaderTimeout time.Duration
|
||||
MaxRetries int
|
||||
RetryBackoff time.Duration
|
||||
CircuitThreshold int
|
||||
CircuitOpenDuration time.Duration
|
||||
}
|
||||
|
||||
func DefaultResiliencePolicy() ResiliencePolicy {
|
||||
return ResiliencePolicy{
|
||||
ResponseHeaderTimeout: 60 * time.Second, MaxRetries: 2, RetryBackoff: 50 * time.Millisecond,
|
||||
CircuitThreshold: 5, CircuitOpenDuration: 30 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
type circuitBreaker struct {
|
||||
mu sync.Mutex
|
||||
failures int
|
||||
threshold int
|
||||
openFor time.Duration
|
||||
openUntil time.Time
|
||||
halfOpenRun bool
|
||||
}
|
||||
|
||||
func newCircuitBreaker(policy ...ResiliencePolicy) *circuitBreaker {
|
||||
settings := DefaultResiliencePolicy()
|
||||
if len(policy) > 0 {
|
||||
settings = policy[0]
|
||||
}
|
||||
return &circuitBreaker{threshold: settings.CircuitThreshold, openFor: settings.CircuitOpenDuration}
|
||||
}
|
||||
|
||||
func (c *circuitBreaker) allow(now time.Time) bool {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.openUntil.IsZero() {
|
||||
return true
|
||||
}
|
||||
if now.Before(c.openUntil) || c.halfOpenRun {
|
||||
return false
|
||||
}
|
||||
c.halfOpenRun = true
|
||||
return true
|
||||
}
|
||||
|
||||
func (c *circuitBreaker) success() {
|
||||
c.mu.Lock()
|
||||
c.failures = 0
|
||||
c.openUntil = time.Time{}
|
||||
c.halfOpenRun = false
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
func (c *circuitBreaker) failure(now time.Time) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.halfOpenRun = false
|
||||
c.failures++
|
||||
if c.failures >= c.threshold || !c.openUntil.IsZero() {
|
||||
c.openUntil = now.Add(c.openFor)
|
||||
}
|
||||
}
|
||||
|
||||
type resilientTransport struct {
|
||||
base http.RoundTripper
|
||||
circuit *circuitBreaker
|
||||
maxRetries int
|
||||
backoff time.Duration
|
||||
}
|
||||
|
||||
func (t *resilientTransport) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
if !t.circuit.allow(time.Now()) {
|
||||
return nil, ErrCircuitOpen
|
||||
}
|
||||
replayable := request.Method == http.MethodGet || request.Method == http.MethodHead ||
|
||||
request.Header.Get("Idempotency-Key") != "" && request.GetBody != nil
|
||||
attempts := 1
|
||||
if replayable {
|
||||
attempts += t.maxRetries
|
||||
}
|
||||
var response *http.Response
|
||||
var err error
|
||||
for attempt := 0; attempt < attempts; attempt++ {
|
||||
current := request
|
||||
if attempt > 0 {
|
||||
if waitErr := waitBackoff(request.Context(), t.backoff*time.Duration(attempt)); waitErr != nil {
|
||||
err = waitErr
|
||||
break
|
||||
}
|
||||
current = request.Clone(request.Context())
|
||||
if request.Body != nil && request.Body != http.NoBody {
|
||||
body, bodyErr := request.GetBody()
|
||||
if bodyErr != nil {
|
||||
err = bodyErr
|
||||
break
|
||||
}
|
||||
current.Body = body
|
||||
}
|
||||
}
|
||||
response, err = t.base.RoundTrip(current)
|
||||
if !retryableResult(response, err) || attempt == attempts-1 {
|
||||
break
|
||||
}
|
||||
if response != nil {
|
||||
_, _ = io.CopyN(io.Discard, response.Body, 4096)
|
||||
_ = response.Body.Close()
|
||||
}
|
||||
}
|
||||
if retryableResult(response, err) {
|
||||
t.circuit.failure(time.Now())
|
||||
} else {
|
||||
t.circuit.success()
|
||||
}
|
||||
return response, err
|
||||
}
|
||||
|
||||
func retryableResult(response *http.Response, err error) bool {
|
||||
if err != nil {
|
||||
return !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded)
|
||||
}
|
||||
return response != nil && (response.StatusCode == http.StatusBadGateway || response.StatusCode == http.StatusServiceUnavailable || response.StatusCode == http.StatusGatewayTimeout)
|
||||
}
|
||||
|
||||
func waitBackoff(ctx context.Context, duration time.Duration) error {
|
||||
if duration <= 0 {
|
||||
return nil
|
||||
}
|
||||
timer := time.NewTimer(duration)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { return f(request) }
|
||||
|
||||
func TestResilientTransportRetriesReplayableRequest(t *testing.T) {
|
||||
attempts := 0
|
||||
transport := &resilientTransport{base: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
attempts++
|
||||
if attempts < 3 {
|
||||
return &http.Response{StatusCode: http.StatusServiceUnavailable, Body: io.NopCloser(strings.NewReader("busy")), Header: make(http.Header)}, nil
|
||||
}
|
||||
return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("ok")), Header: make(http.Header)}, nil
|
||||
}), circuit: newCircuitBreaker(), maxRetries: 2}
|
||||
request, _ := http.NewRequest(http.MethodGet, "https://provider.example/v1/models", nil)
|
||||
response, err := transport.RoundTrip(request)
|
||||
if err != nil || response.StatusCode != http.StatusOK || attempts != 3 {
|
||||
t.Fatalf("response=%v err=%v attempts=%d", response, err, attempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResilientTransportDoesNotRetryUnsafePost(t *testing.T) {
|
||||
attempts := 0
|
||||
transport := &resilientTransport{base: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
attempts++
|
||||
return nil, errors.New("network failure")
|
||||
}), circuit: newCircuitBreaker(), maxRetries: 2}
|
||||
request, _ := http.NewRequest(http.MethodPost, "https://provider.example/v1/chat/completions", strings.NewReader("{}"))
|
||||
_, _ = transport.RoundTrip(request)
|
||||
if attempts != 1 {
|
||||
t.Fatalf("unsafe POST attempted %d times", attempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCircuitOpensAndAllowsSingleProbe(t *testing.T) {
|
||||
circuit := newCircuitBreaker()
|
||||
circuit.threshold = 2
|
||||
circuit.openFor = time.Millisecond
|
||||
now := time.Now()
|
||||
circuit.failure(now)
|
||||
circuit.failure(now)
|
||||
if circuit.allow(now) {
|
||||
t.Fatal("open circuit allowed request")
|
||||
}
|
||||
if !circuit.allow(now.Add(2 * time.Millisecond)) {
|
||||
t.Fatal("circuit did not allow half-open probe")
|
||||
}
|
||||
if circuit.allow(now.Add(2 * time.Millisecond)) {
|
||||
t.Fatal("circuit allowed concurrent half-open probe")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
type ModelRouteQuery struct {
|
||||
ProviderCode string
|
||||
Model string
|
||||
Endpoint string
|
||||
APIKeyID string
|
||||
TenantID string
|
||||
Seed string
|
||||
}
|
||||
|
||||
type ModelRouteResult struct {
|
||||
ResolvedAdapter
|
||||
TargetModel string
|
||||
Matched bool
|
||||
Known bool
|
||||
}
|
||||
|
||||
type ModelRouteResolver interface {
|
||||
ModelRoutingEnabled() bool
|
||||
ResolveModelRoute(ModelRouteQuery) (ModelRouteResult, error)
|
||||
}
|
||||
|
||||
type modelRequest struct {
|
||||
body map[string]json.RawMessage
|
||||
model string
|
||||
}
|
||||
|
||||
func readModelRequest(request *http.Request, maxBody int64) (modelRequest, error) {
|
||||
if request.Body == nil || request.Method == http.MethodGet {
|
||||
return modelRequest{}, nil
|
||||
}
|
||||
body, err := io.ReadAll(io.LimitReader(request.Body, maxBody+1))
|
||||
if err != nil {
|
||||
return modelRequest{}, err
|
||||
}
|
||||
_ = request.Body.Close()
|
||||
if int64(len(body)) > maxBody {
|
||||
return modelRequest{}, errRequestBodyTooLarge
|
||||
}
|
||||
restoreRequestBody(request, body)
|
||||
var payload map[string]json.RawMessage
|
||||
if json.Unmarshal(body, &payload) != nil {
|
||||
return modelRequest{}, nil
|
||||
}
|
||||
var model string
|
||||
_ = json.Unmarshal(payload["model"], &model)
|
||||
return modelRequest{body: payload, model: model}, nil
|
||||
}
|
||||
|
||||
func (m modelRequest) rewrite(request *http.Request, target string) error {
|
||||
if len(m.body) == 0 || target == "" || target == m.model {
|
||||
return nil
|
||||
}
|
||||
encodedModel, _ := json.Marshal(target)
|
||||
m.body["model"] = encodedModel
|
||||
body, err := json.Marshal(m.body)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
restoreRequestBody(request, body)
|
||||
return nil
|
||||
}
|
||||
|
||||
func restoreRequestBody(request *http.Request, body []byte) {
|
||||
request.Body = io.NopCloser(bytes.NewReader(body))
|
||||
request.ContentLength = int64(len(body))
|
||||
request.GetBody = func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewReader(body)), nil }
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/apikey"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
var ErrTokenQuotaUnavailable = errors.New("token quota unavailable")
|
||||
|
||||
type TokenReservation struct {
|
||||
Allowed bool
|
||||
APIKeyID string
|
||||
CounterKey string
|
||||
Reserved int64
|
||||
Limit int64
|
||||
Remaining int64
|
||||
ResetAt time.Time
|
||||
}
|
||||
|
||||
type TokenQuotaController interface {
|
||||
Reserve(context.Context, apikey.Principal, int64, time.Time) (TokenReservation, error)
|
||||
Commit(context.Context, TokenReservation, int64) error
|
||||
}
|
||||
|
||||
type RedisTokenQuotaController struct {
|
||||
client *redis.Client
|
||||
reserveScript *redis.Script
|
||||
commitScript *redis.Script
|
||||
}
|
||||
|
||||
func NewRedisTokenQuotaController(client *redis.Client) *RedisTokenQuotaController {
|
||||
return &RedisTokenQuotaController{
|
||||
client: client, reserveScript: redis.NewScript(tokenReserveScript), commitScript: redis.NewScript(tokenCommitScript),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *RedisTokenQuotaController) Reserve(ctx context.Context, principal apikey.Principal, estimate int64, now time.Time) (TokenReservation, error) {
|
||||
if principal.APIKeyID == "" {
|
||||
return TokenReservation{Allowed: true}, nil
|
||||
}
|
||||
if estimate < 0 {
|
||||
estimate = 0
|
||||
}
|
||||
// Accounts without a monthly token quota never need a reservation and must
|
||||
// not create a pointless monthly counter key in Redis.
|
||||
if principal.MonthlyTokenQuota == 0 {
|
||||
return TokenReservation{Allowed: true}, nil
|
||||
}
|
||||
if c == nil || c.client == nil {
|
||||
return TokenReservation{}, ErrTokenQuotaUnavailable
|
||||
}
|
||||
now = now.UTC()
|
||||
reset := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, time.UTC)
|
||||
key := apikey.MonthlyTokenUsageKey(principal.APIKeyID, now)
|
||||
result, err := c.reserveScript.Run(ctx, c.client, []string{key}, estimate, principal.MonthlyTokenQuota, int64(reset.Sub(now).Seconds())+86400).Slice()
|
||||
if err != nil || len(result) != 2 {
|
||||
if principal.MonthlyTokenQuota == 0 {
|
||||
return TokenReservation{Allowed: true}, nil
|
||||
}
|
||||
return TokenReservation{}, fmt.Errorf("%w: %v", ErrTokenQuotaUnavailable, err)
|
||||
}
|
||||
allowed, err := redisInteger(result[0])
|
||||
if err != nil {
|
||||
return TokenReservation{}, ErrTokenQuotaUnavailable
|
||||
}
|
||||
current, err := redisInteger(result[1])
|
||||
if err != nil {
|
||||
return TokenReservation{}, ErrTokenQuotaUnavailable
|
||||
}
|
||||
return TokenReservation{
|
||||
Allowed: allowed == 1, APIKeyID: principal.APIKeyID, CounterKey: key, Reserved: estimate,
|
||||
Limit: principal.MonthlyTokenQuota, Remaining: max(principal.MonthlyTokenQuota-current, 0), ResetAt: reset,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *RedisTokenQuotaController) Commit(ctx context.Context, reservation TokenReservation, actual int64) error {
|
||||
if c == nil || c.client == nil || reservation.CounterKey == "" || !reservation.Allowed {
|
||||
return nil
|
||||
}
|
||||
if actual < 0 {
|
||||
actual = 0
|
||||
}
|
||||
if _, err := c.commitScript.Run(ctx, c.client, []string{reservation.CounterKey}, actual-reservation.Reserved).Result(); err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrTokenQuotaUnavailable, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
const tokenReserveScript = `
|
||||
local estimate = tonumber(ARGV[1])
|
||||
local quota = tonumber(ARGV[2])
|
||||
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
|
||||
if quota > 0 and current + estimate > quota then
|
||||
return {0, current}
|
||||
end
|
||||
if estimate > 0 then
|
||||
current = redis.call('INCRBY', KEYS[1], estimate)
|
||||
if current == estimate then redis.call('EXPIRE', KEYS[1], tonumber(ARGV[3])) end
|
||||
elseif redis.call('EXISTS', KEYS[1]) == 0 then
|
||||
redis.call('SET', KEYS[1], 0, 'EX', tonumber(ARGV[3]))
|
||||
end
|
||||
return {1, current}
|
||||
`
|
||||
|
||||
const tokenCommitScript = `
|
||||
local delta = tonumber(ARGV[1])
|
||||
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
|
||||
local updated = current + delta
|
||||
if updated < 0 then updated = 0 end
|
||||
redis.call('SET', KEYS[1], updated, 'KEEPTTL')
|
||||
return updated
|
||||
`
|
||||
|
||||
func writeTokenQuotaHeaders(header mapHeader, reservation TokenReservation) {
|
||||
if reservation.Limit <= 0 || reservation.ResetAt.IsZero() {
|
||||
return
|
||||
}
|
||||
header.Set("X-TokenLimit-Limit", strconv.FormatInt(reservation.Limit, 10))
|
||||
header.Set("X-TokenLimit-Remaining", strconv.FormatInt(max(reservation.Remaining, 0), 10))
|
||||
header.Set("X-TokenLimit-Reset", strconv.FormatInt(reservation.ResetAt.Unix(), 10))
|
||||
}
|
||||
|
||||
type mapHeader interface{ Set(string, string) }
|
||||
@@ -0,0 +1,261 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log/slog"
|
||||
"mime"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// maxUsageDocumentBytes bounds the buffered tail of a non-streaming response.
|
||||
// The window is kept from the end of the body (where providers place the
|
||||
// "usage" object), so token accounting stays correct for responses far larger
|
||||
// than this bound while memory use stays bounded per in-flight response.
|
||||
const maxUsageDocumentBytes = 2 << 20
|
||||
|
||||
type usageSession struct {
|
||||
controller TokenQuotaController
|
||||
reservation TokenReservation
|
||||
fallback int64
|
||||
onFinish func(TokenUsage)
|
||||
once sync.Once
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
type TokenUsage struct {
|
||||
Input int64
|
||||
Output int64
|
||||
Total int64
|
||||
}
|
||||
|
||||
type usageSessionContextKey struct{}
|
||||
|
||||
func withUsageSession(request *http.Request, session *usageSession) *http.Request {
|
||||
return request.WithContext(context.WithValue(request.Context(), usageSessionContextKey{}, session))
|
||||
}
|
||||
|
||||
func usageSessionFrom(request *http.Request) *usageSession {
|
||||
session, _ := request.Context().Value(usageSessionContextKey{}).(*usageSession)
|
||||
return session
|
||||
}
|
||||
|
||||
func (s *usageSession) finish(usage TokenUsage) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.once.Do(func() {
|
||||
usage.Total = max(usage.Total, usage.Input+usage.Output)
|
||||
if usage.Total <= 0 {
|
||||
usage.Total = s.fallback
|
||||
}
|
||||
if s.controller != nil {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
if err := s.controller.Commit(ctx, s.reservation, usage.Total); err != nil && s.log != nil {
|
||||
// The reservation already counted towards the quota, so a
|
||||
// failed reconciliation silently leaves the counter slightly
|
||||
// off. Surface it instead of dropping it.
|
||||
s.log.Warn("token quota commit failed", "error", err)
|
||||
}
|
||||
cancel()
|
||||
}
|
||||
if s.onFinish != nil {
|
||||
s.onFinish(usage)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
type usageReadCloser struct {
|
||||
io.ReadCloser
|
||||
collector *usageCollector
|
||||
session *usageSession
|
||||
}
|
||||
|
||||
func newUsageReadCloser(body io.ReadCloser, contentType string, session *usageSession) io.ReadCloser {
|
||||
mediaType, _, _ := mime.ParseMediaType(contentType)
|
||||
return &usageReadCloser{ReadCloser: body, collector: &usageCollector{sse: mediaType == "text/event-stream"}, session: session}
|
||||
}
|
||||
|
||||
func (r *usageReadCloser) Read(buffer []byte) (int, error) {
|
||||
n, err := r.ReadCloser.Read(buffer)
|
||||
if n > 0 {
|
||||
r.collector.feed(buffer[:n])
|
||||
}
|
||||
if err == io.EOF {
|
||||
r.session.finish(r.collector.usage())
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (r *usageReadCloser) Close() error {
|
||||
r.session.finish(r.collector.usage())
|
||||
return r.ReadCloser.Close()
|
||||
}
|
||||
|
||||
type usageCollector struct {
|
||||
sse bool
|
||||
pending []byte
|
||||
doc []byte
|
||||
input int64
|
||||
output int64
|
||||
total int64
|
||||
}
|
||||
|
||||
func (c *usageCollector) feed(chunk []byte) {
|
||||
if !c.sse {
|
||||
c.doc = append(c.doc, chunk...)
|
||||
// Keep only a bounded tail window. The "usage" member lives at the end
|
||||
// of a non-streaming response, so dropping the head (never the tail)
|
||||
// preserves accounting for arbitrarily large bodies at a fixed memory
|
||||
// cost instead of truncating usage away.
|
||||
if len(c.doc) > maxUsageDocumentBytes {
|
||||
c.doc = append([]byte(nil), c.doc[len(c.doc)-maxUsageDocumentBytes:]...)
|
||||
}
|
||||
return
|
||||
}
|
||||
c.pending = append(c.pending, chunk...)
|
||||
for {
|
||||
index := bytes.IndexByte(c.pending, '\n')
|
||||
if index < 0 {
|
||||
if len(c.pending) > maxUsageDocumentBytes {
|
||||
c.pending = c.pending[:0]
|
||||
}
|
||||
return
|
||||
}
|
||||
line := strings.TrimSpace(string(c.pending[:index]))
|
||||
c.pending = c.pending[index+1:]
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if payload != "" && payload != "[DONE]" {
|
||||
c.consumeJSON([]byte(payload))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *usageCollector) tokens() int64 {
|
||||
return c.usage().Total
|
||||
}
|
||||
|
||||
func (c *usageCollector) usage() TokenUsage {
|
||||
if c.sse && len(c.pending) > 0 {
|
||||
line := strings.TrimSpace(string(c.pending))
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
c.consumeJSON([]byte(strings.TrimSpace(strings.TrimPrefix(line, "data:"))))
|
||||
}
|
||||
c.pending = nil
|
||||
}
|
||||
if !c.sse && len(c.doc) > 0 {
|
||||
c.consumeJSON(c.doc) // fast path: whole valid JSON object
|
||||
c.consumeUsageObject(c.doc) // tail extraction: covers truncated bodies
|
||||
}
|
||||
return TokenUsage{Input: c.input, Output: c.output, Total: max(c.total, c.input+c.output)}
|
||||
}
|
||||
|
||||
// consumeUsageObject extracts the trailing `"usage"` object from a possibly
|
||||
// truncated response document. It scans for the final `"usage"` member and
|
||||
// decodes the JSON object that immediately follows — where OpenAI-compatible
|
||||
// providers place token accounting in non-streaming responses — so usage is
|
||||
// still counted when the buffered window starts mid-object.
|
||||
func (c *usageCollector) consumeUsageObject(doc []byte) {
|
||||
const key = `"usage"`
|
||||
index := bytes.LastIndex(doc, []byte(key))
|
||||
if index < 0 {
|
||||
return
|
||||
}
|
||||
rest := doc[index+len(key):]
|
||||
colon := bytes.IndexByte(rest, ':')
|
||||
if colon < 0 {
|
||||
return
|
||||
}
|
||||
rest = bytes.TrimSpace(rest[colon+1:])
|
||||
if len(rest) == 0 || rest[0] != '{' {
|
||||
return
|
||||
}
|
||||
depth := 0
|
||||
end := -1
|
||||
inString := false
|
||||
escaped := false
|
||||
for i := 0; i < len(rest); i++ {
|
||||
ch := rest[i]
|
||||
if inString {
|
||||
if escaped {
|
||||
escaped = false
|
||||
} else if ch == '\\' {
|
||||
escaped = true
|
||||
} else if ch == '"' {
|
||||
inString = false
|
||||
}
|
||||
continue
|
||||
}
|
||||
switch ch {
|
||||
case '"':
|
||||
inString = true
|
||||
case '{':
|
||||
depth++
|
||||
case '}':
|
||||
depth--
|
||||
if depth == 0 {
|
||||
end = i + 1
|
||||
}
|
||||
}
|
||||
if end > 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
if end > 0 {
|
||||
c.consumeJSON(rest[:end])
|
||||
}
|
||||
}
|
||||
|
||||
func (c *usageCollector) consumeJSON(payload []byte) {
|
||||
decoder := json.NewDecoder(bytes.NewReader(payload))
|
||||
decoder.UseNumber()
|
||||
var value any
|
||||
if decoder.Decode(&value) == nil {
|
||||
c.walk(value, false)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *usageCollector) walk(value any, inUsage bool) {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
for key, child := range typed {
|
||||
usage := inUsage || key == "usage"
|
||||
if usage {
|
||||
switch key {
|
||||
case "input_tokens", "prompt_tokens":
|
||||
c.input = max(c.input, jsonInt64(child))
|
||||
case "output_tokens", "completion_tokens":
|
||||
c.output = max(c.output, jsonInt64(child))
|
||||
case "total_tokens":
|
||||
c.total = max(c.total, jsonInt64(child))
|
||||
}
|
||||
}
|
||||
c.walk(child, usage)
|
||||
}
|
||||
case []any:
|
||||
for _, child := range typed {
|
||||
c.walk(child, inUsage)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func jsonInt64(value any) int64 {
|
||||
switch number := value.(type) {
|
||||
case json.Number:
|
||||
parsed, _ := number.Int64()
|
||||
return max(parsed, 0)
|
||||
case float64:
|
||||
return max(int64(number), 0)
|
||||
case int64:
|
||||
return max(number, 0)
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
package gateway
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestUsageCollectorReadsOpenAIJSON(t *testing.T) {
|
||||
collector := &usageCollector{}
|
||||
collector.feed([]byte(`{"usage":{"prompt_tokens":11,"completion_tokens":7,"total_tokens":18}}`))
|
||||
if got := collector.tokens(); got != 18 {
|
||||
t.Fatalf("got %d, want 18", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageCollectorReadsChunkedSSEAndAnthropicUsage(t *testing.T) {
|
||||
collector := &usageCollector{sse: true}
|
||||
collector.feed([]byte("event: message_start\ndata: {\"message\":{\"usage\":{\"input_tokens\":13}}}\n\n"))
|
||||
collector.feed([]byte("event: message_delta\ndata: {\"usage\":{\"output_"))
|
||||
collector.feed([]byte("tokens\":9}}\n\ndata: [DONE]\n\n"))
|
||||
if got := collector.tokens(); got != 22 {
|
||||
t.Fatalf("got %d, want 22", got)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user