openteam/server/internal/proxy/gateway.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")
}