SakurasanandClaude Sonnet 5
d0e31b198f
Gin + GORM + pure-Go SQLite. Users/auth (JWT), API key management with quotas, proxy gateway with weighted channel failover and health checks, usage/billing ledger, cross-protocol conversion (Anthropic Messages / OpenAI Chat Completions / OpenAI Responses), and channel/model admin API. Channels declare native API formats and auto-convert the rest. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
415 lines
14 KiB
Go
415 lines
14 KiB
Go
package proxy
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/google/uuid"
|
|
"go.uber.org/zap"
|
|
"gorm.io/gorm"
|
|
|
|
"openteam/server/internal/apikey"
|
|
"openteam/server/internal/billing"
|
|
"openteam/server/internal/channel"
|
|
"openteam/server/internal/config"
|
|
"openteam/server/internal/pkg/ratelimit"
|
|
"openteam/server/internal/usage"
|
|
"openteam/server/internal/store"
|
|
)
|
|
|
|
// Protocol identifies the client-facing API protocol.
|
|
type Protocol string
|
|
|
|
const (
|
|
ProtocolOpenAIChat Protocol = "openai-chat"
|
|
ProtocolOpenAIResponses Protocol = "openai-responses"
|
|
ProtocolAnthropic Protocol = "anthropic"
|
|
)
|
|
|
|
// Route describes one proxied endpoint.
|
|
type Route struct {
|
|
Protocol Protocol
|
|
UpstreamPath string // suffix after base URL, e.g. /v1/chat/completions
|
|
NativeProvider string // provider type that matches this protocol ("openai" | "anthropic")
|
|
}
|
|
|
|
var Routes = []Route{
|
|
{Protocol: ProtocolOpenAIChat, UpstreamPath: "/v1/chat/completions", NativeProvider: "openai"},
|
|
{Protocol: ProtocolOpenAIResponses, UpstreamPath: "/v1/responses", NativeProvider: "openai"},
|
|
{Protocol: ProtocolAnthropic, UpstreamPath: "/v1/messages", NativeProvider: "anthropic"},
|
|
}
|
|
|
|
const maxBodyBytes = 16 << 20 // 16 MiB
|
|
|
|
type Gateway struct {
|
|
db *gorm.DB
|
|
cfg *config.Config
|
|
log *zap.Logger
|
|
channel *channel.Service
|
|
billing *billing.Service
|
|
usage *usage.Service
|
|
apiKeys *apikey.Service
|
|
limiter *ratelimit.Limiter
|
|
client *http.Client
|
|
}
|
|
|
|
func NewGateway(db *gorm.DB, cfg *config.Config, log *zap.Logger,
|
|
ch *channel.Service, bill *billing.Service, use *usage.Service, ak *apikey.Service) *Gateway {
|
|
return &Gateway{
|
|
db: db, cfg: cfg, log: log,
|
|
channel: ch, billing: bill, usage: use, apiKeys: ak,
|
|
limiter: ratelimit.New(float64(cfg.RateLimit.RequestsPerMin)/60.0, cfg.RateLimit.Burst),
|
|
client: &http.Client{
|
|
// Transport-level timeout; stream reads rely on the context so a
|
|
// connected-but-silent upstream is still bounded.
|
|
Timeout: time.Duration(cfg.Proxy.DefaultTimeoutMs) * time.Millisecond,
|
|
},
|
|
}
|
|
}
|
|
|
|
// authContext carries the authenticated key + user through a proxy request.
|
|
type authContext struct {
|
|
key *store.ApiKey
|
|
user *store.User
|
|
}
|
|
|
|
// Handle builds a gin handler for a route.
|
|
func (g *Gateway) Handle(route Route) gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
g.proxy(c, route)
|
|
}
|
|
}
|
|
|
|
func (g *Gateway) proxy(c *gin.Context, route Route) {
|
|
start := time.Now()
|
|
reqID := uuid.NewString()
|
|
c.Header("X-Request-Id", reqID)
|
|
|
|
auth, err := g.authenticate(c)
|
|
if err != nil {
|
|
g.writeProxyError(c, route.Protocol, http.StatusUnauthorized, "invalid API key")
|
|
return
|
|
}
|
|
|
|
body, err := io.ReadAll(io.LimitReader(c.Request.Body, maxBodyBytes))
|
|
if err != nil {
|
|
g.writeProxyError(c, route.Protocol, http.StatusBadRequest, "failed to read request body")
|
|
return
|
|
}
|
|
if len(body) == 0 {
|
|
g.writeProxyError(c, route.Protocol, http.StatusBadRequest, "empty request body")
|
|
return
|
|
}
|
|
|
|
var meta struct {
|
|
Model string `json:"model"`
|
|
Stream bool `json:"stream"`
|
|
}
|
|
if err := json.Unmarshal(body, &meta); err != nil {
|
|
g.writeProxyError(c, route.Protocol, http.StatusBadRequest, "invalid JSON body: "+err.Error())
|
|
return
|
|
}
|
|
if meta.Model == "" {
|
|
g.writeProxyError(c, route.Protocol, http.StatusBadRequest, "missing model field")
|
|
return
|
|
}
|
|
|
|
// Per-key model whitelist.
|
|
if len(auth.key.AllowedModels) > 0 && !contains(auth.key.AllowedModels, meta.Model) {
|
|
g.writeProxyError(c, route.Protocol, http.StatusForbidden,
|
|
"model not allowed for this API key: "+meta.Model)
|
|
return
|
|
}
|
|
|
|
// Rate limit (user + key).
|
|
if !g.limiter.Allow(fmt.Sprintf("u:%d", auth.user.ID)) {
|
|
g.writeProxyError(c, route.Protocol, http.StatusTooManyRequests, "rate limit exceeded")
|
|
return
|
|
}
|
|
if !g.limiter.Allow(fmt.Sprintf("k:%d", auth.key.ID)) {
|
|
g.writeProxyError(c, route.Protocol, http.StatusTooManyRequests, "key rate limit exceeded")
|
|
return
|
|
}
|
|
|
|
// Resolve model.
|
|
model, err := g.channel.ResolveModel(meta.Model)
|
|
if err != nil {
|
|
g.writeProxyError(c, route.Protocol, http.StatusNotFound, err.Error())
|
|
return
|
|
}
|
|
|
|
// Channel attempts with failover: on transport/5xx failures the request is
|
|
// retried against another channel bound to the same model.
|
|
exclude := map[int64]bool{}
|
|
attempts := 1 + g.cfg.Proxy.MaxRetries
|
|
if attempts < 1 {
|
|
attempts = 1
|
|
}
|
|
var lastCh *store.Channel
|
|
formatBlocked := false
|
|
attempted := false
|
|
for attempt := 0; attempt < attempts; attempt++ {
|
|
ch, binding, err := g.channel.SelectChannel(model.ID, exclude)
|
|
if err != nil {
|
|
// Every bound channel was skipped for format reasons: say so
|
|
// clearly instead of reporting a generic upstream failure.
|
|
if formatBlocked && !attempted {
|
|
g.writeProxyError(c, route.Protocol, http.StatusBadRequest,
|
|
"no channel supports the "+string(route.Protocol)+" API format")
|
|
return
|
|
}
|
|
if attempt == 0 {
|
|
g.writeProxyError(c, route.Protocol, http.StatusServiceUnavailable, "no available channel for model")
|
|
} else {
|
|
g.recordError(reqID, auth, model, lastCh, start, "upstream_error")
|
|
g.writeProxyError(c, route.Protocol, http.StatusBadGateway, "all upstream channels failed")
|
|
}
|
|
return
|
|
}
|
|
exclude[ch.ID] = true
|
|
lastCh = ch
|
|
|
|
// Skip channels that cannot serve this protocol at all (no native
|
|
// support and no conversion path).
|
|
if !channelServesFormat(ch, route) {
|
|
formatBlocked = true
|
|
g.log.Debug("channel cannot serve protocol",
|
|
zap.String("protocol", string(route.Protocol)), zap.Int64("channel_id", ch.ID))
|
|
continue
|
|
}
|
|
|
|
// Optional balance pre-check with an estimate.
|
|
if g.cfg.Proxy.BillingExactBalance {
|
|
estIn := int64(len(body) / 4)
|
|
estOut := int64(512)
|
|
if meta.Stream {
|
|
estOut = int64(g.cfg.Proxy.DefaultMaxTokens)
|
|
}
|
|
est, cerr := g.billing.EstimateCost(model.ID, estIn, estOut, 0)
|
|
if cerr == nil {
|
|
if berr := g.billing.CheckBalance(auth.user.ID, est); berr != nil {
|
|
g.recordError(reqID, auth, model, ch, start, "insufficient_balance")
|
|
g.writeProxyError(c, route.Protocol, http.StatusPaymentRequired, "insufficient balance")
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
upstreamKey, err := g.channel.DecryptKey(ch.APIKeyEnc)
|
|
if err != nil {
|
|
g.log.Error("decrypt channel key", zap.Error(err), zap.Int64("channel_id", ch.ID))
|
|
g.writeProxyError(c, route.Protocol, http.StatusBadGateway, "channel key unavailable")
|
|
return
|
|
}
|
|
|
|
release, err := g.channel.Acquire(ch.ID)
|
|
if err != nil {
|
|
g.writeProxyError(c, route.Protocol, http.StatusServiceUnavailable, "channel unavailable")
|
|
return
|
|
}
|
|
attempted = true
|
|
retryable := g.forward(c, route, meta.Model, ch, binding.UpstreamModel, upstreamKey, body, start, reqID, auth, model)
|
|
release()
|
|
if !retryable {
|
|
return
|
|
}
|
|
g.log.Warn("upstream failed, retrying on another channel",
|
|
zap.Int64("model_id", model.ID), zap.Int64("channel_id", ch.ID), zap.Int("attempt", attempt+1))
|
|
}
|
|
if formatBlocked && !attempted {
|
|
g.writeProxyError(c, route.Protocol, http.StatusBadRequest,
|
|
"no channel supports the "+string(route.Protocol)+" API format")
|
|
return
|
|
}
|
|
g.recordError(reqID, auth, model, lastCh, start, "upstream_error")
|
|
g.writeProxyError(c, route.Protocol, http.StatusBadGateway, "all upstream channels failed")
|
|
}
|
|
|
|
// forward decides passthrough vs conversion and calls the upstream.
|
|
// It returns true when the failure is retryable on another channel.
|
|
func (g *Gateway) forward(c *gin.Context, route Route, clientModel string,
|
|
ch *store.Channel, upstreamModel, upstreamKey string, body []byte,
|
|
start time.Time, reqID string, auth *authContext, model *store.Model) bool {
|
|
|
|
upstreamBody, converted := g.prepareUpstreamBody(route, ch, body, upstreamModel)
|
|
if converted {
|
|
c.Header("x-converted", "true")
|
|
}
|
|
|
|
// Build upstream request bound to the client context so disconnects cancel it.
|
|
ctx := c.Request.Context()
|
|
upstreamURL := strings.TrimSuffix(ch.BaseURL, "/") + upstreamPath(route, ch)
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(upstreamBody))
|
|
if err != nil {
|
|
g.writeProxyError(c, route.Protocol, http.StatusBadGateway, "failed to build upstream request")
|
|
return false
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Authorization", "Bearer "+upstreamKey)
|
|
req.Header.Set("Accept", "application/json")
|
|
req.Header.Set("X-Request-Id", reqID)
|
|
// Explicitly drop hop-by-hop / auth-ish headers we don't want forwarded.
|
|
copyProxyHeaders(c, req)
|
|
|
|
resp, err := g.client.Do(req)
|
|
if err != nil {
|
|
// Client disconnect vs upstream failure.
|
|
if errors.Is(ctx.Err(), context.Canceled) {
|
|
g.recordCanceled(reqID, auth, model, ch, start)
|
|
return false
|
|
}
|
|
g.log.Warn("upstream request failed", zap.Error(err), zap.Int64("channel_id", ch.ID))
|
|
return true
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode >= 500 {
|
|
errBody, _ := io.ReadAll(resp.Body)
|
|
g.log.Warn("upstream returned 5xx", zap.Int("status", resp.StatusCode),
|
|
zap.Int64("channel_id", ch.ID), zap.String("body", truncateText(string(errBody), 512)))
|
|
return true
|
|
}
|
|
if resp.StatusCode >= 400 {
|
|
errBody, _ := io.ReadAll(resp.Body)
|
|
g.recordError(reqID, auth, model, ch, start, "upstream_"+strconv.Itoa(resp.StatusCode))
|
|
g.writeUpstreamError(c, route.Protocol, resp.StatusCode, errBody)
|
|
return false
|
|
}
|
|
|
|
streaming := bodyStreamFlag(body, route)
|
|
if streaming {
|
|
g.streamResponse(c, route, resp, start, reqID, auth, model, ch)
|
|
} else {
|
|
g.plainResponse(c, route, resp, start, reqID, auth, model, ch)
|
|
}
|
|
return false
|
|
}
|
|
|
|
// prepareUpstreamBody rewrites the model name, or converts the body when the
|
|
// channel does not serve the client protocol natively. Returns the payload and
|
|
// whether any conversion happened.
|
|
func (g *Gateway) prepareUpstreamBody(route Route, ch *store.Channel, body []byte, upstreamModel string) ([]byte, bool) {
|
|
if ch.SupportsFormat(string(route.Protocol)) {
|
|
if upstreamModel == "" || sameModel(body, upstreamModel) {
|
|
return body, false
|
|
}
|
|
rewritten, err := setModelField(body, upstreamModel)
|
|
if err != nil {
|
|
return body, false
|
|
}
|
|
return rewritten, true
|
|
}
|
|
converted, err := convertRequest(route, ch.Provider, body, upstreamModel)
|
|
if err != nil {
|
|
g.log.Warn("request conversion failed, falling back to passthrough",
|
|
zap.Error(err), zap.String("route", string(route.Protocol)), zap.String("provider", ch.Provider))
|
|
return body, false
|
|
}
|
|
return converted, true
|
|
}
|
|
|
|
// streamResponse forwards an SSE stream to the client while extracting usage.
|
|
func (g *Gateway) streamResponse(c *gin.Context, route Route, resp *http.Response,
|
|
start time.Time, reqID string, auth *authContext, model *store.Model, ch *store.Channel) {
|
|
|
|
c.Header("Content-Type", "text/event-stream")
|
|
c.Header("Cache-Control", "no-cache")
|
|
c.Header("Connection", "keep-alive")
|
|
|
|
if upstreamProvider(route, ch) == "anthropic" {
|
|
g.streamAnthropic(c, route, resp, start, reqID, auth, model, ch)
|
|
return
|
|
}
|
|
g.streamOpenAI(c, route, resp, start, reqID, auth, model, ch)
|
|
}
|
|
|
|
// plainResponse buffers a non-stream upstream response and returns it.
|
|
func (g *Gateway) plainResponse(c *gin.Context, route Route, resp *http.Response,
|
|
start time.Time, reqID string, auth *authContext, model *store.Model, ch *store.Channel) {
|
|
|
|
raw, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
g.log.Warn("read upstream body", zap.Error(err))
|
|
g.writeProxyError(c, route.Protocol, http.StatusBadGateway, "failed to read upstream response")
|
|
return
|
|
}
|
|
|
|
out := raw
|
|
usageInfo := parseUsageForProtocol(route.Protocol, raw)
|
|
if !ch.SupportsFormat(string(route.Protocol)) {
|
|
converted, cerr := convertResponse(route, ch.Provider, raw)
|
|
if cerr == nil {
|
|
out = converted
|
|
usageInfo = parseUsageForProtocol(route.Protocol, out)
|
|
c.Header("x-converted", "true")
|
|
} else {
|
|
g.log.Warn("response conversion failed, forwarding raw",
|
|
zap.Error(cerr), zap.String("protocol", string(route.Protocol)))
|
|
}
|
|
}
|
|
c.Data(http.StatusOK, "application/json", out)
|
|
g.afterComplete(start, reqID, auth, model, ch, usageInfo, "success", "")
|
|
}
|
|
|
|
// afterComplete performs billing and usage accounting for a finished request.
|
|
func (g *Gateway) afterComplete(start time.Time, reqID string, auth *authContext,
|
|
model *store.Model, ch *store.Channel, usageInfo *tokenUsage, status, errCode string) {
|
|
|
|
latency := int(time.Since(start).Milliseconds())
|
|
if usageInfo == nil {
|
|
usageInfo = &tokenUsage{}
|
|
}
|
|
cost := billing.CostFromPrices(usageInfo.input, usageInfo.output, usageInfo.cacheRead,
|
|
billing.PriceSnapshot{
|
|
InputPrice: model.InputPrice, OutputPrice: model.OutputPrice, CacheReadPrice: model.CacheReadPrice,
|
|
})
|
|
|
|
go func() {
|
|
// Asynchronous: deduct balance first, then record usage.
|
|
if cost.IsPositive() {
|
|
if _, err := g.billing.Deduct(auth.user.ID, cost, "usage", reqID); err != nil {
|
|
g.log.Warn("deduct balance failed", zap.Error(err),
|
|
zap.Int64("user_id", auth.user.ID), zap.String("request_id", reqID))
|
|
}
|
|
}
|
|
g.usage.Record(usage.Record{
|
|
RequestID: reqID,
|
|
UserID: auth.user.ID,
|
|
KeyID: auth.key.ID,
|
|
ChannelID: ch.ID,
|
|
ModelID: model.ID,
|
|
ModelName: model.Name,
|
|
InputTokens: usageInfo.input,
|
|
OutputTokens: usageInfo.output,
|
|
CacheReadTokens: usageInfo.cacheRead,
|
|
CacheCreationTokens: usageInfo.cacheCreation,
|
|
InputPrice: model.InputPrice,
|
|
OutputPrice: model.OutputPrice,
|
|
CacheReadPrice: model.CacheReadPrice,
|
|
Cost: cost,
|
|
LatencyMs: latency,
|
|
Status: status,
|
|
ErrorCode: errCode,
|
|
})
|
|
g.db.Model(&store.ApiKey{}).Where("id = ?", auth.key.ID).
|
|
Update("last_used_at", time.Now())
|
|
}()
|
|
}
|
|
|
|
func (g *Gateway) recordError(reqID string, auth *authContext, model *store.Model, ch *store.Channel, start time.Time, errCode string) {
|
|
g.afterComplete(start, reqID, auth, model, ch, &tokenUsage{}, "error", errCode)
|
|
}
|
|
|
|
func (g *Gateway) recordCanceled(reqID string, auth *authContext, model *store.Model, ch *store.Channel, start time.Time) {
|
|
g.afterComplete(start, reqID, auth, model, ch, &tokenUsage{}, "canceled", "client_disconnect")
|
|
}
|