openteam/server/internal/proxy/handlers.go

package proxy

import (
	"errors"
	"fmt"
	"net/http"

	"github.com/gin-gonic/gin"
	"github.com/openteam/server/internal/proxy/convert"
	"github.com/openteam/server/internal/store"
)

// failBodyRead 读取请求体失败时的报错：超限单独回 413，让客户端能区分"体太大"和"体畸形"。
func (g *Gateway) failBodyRead(c *gin.Context, err error) {
	var mb *http.MaxBytesError
	if errors.As(err, &mb) {
		apiError(c, http.StatusRequestEntityTooLarge, "request_too_large",
			fmt.Sprintf("请求体超过网关上限 %d MB", g.maxBodyMB))
		return
	}
	apiError(c, http.StatusBadRequest, "invalid_request_error", "failed to read request body")
}

// chatCompletions POST /v1/chat/completions
func (g *Gateway) chatCompletions(c *gin.Context) {
	u, ok := g.resolveUser(c)
	if !ok {
		return
	}
	if !g.checkBalance(c, u) {
		return
	}
	br, body, err := parseBody(c)
	if err != nil {
		g.failBodyRead(c, err)
		return
	}
	c.Set("protocol", convert.ProtoChat)
	c.Set("model_name", br.Model)
	g.recordRawRequest(c, u, body)
	if !g.checkModelExists(c, br.Model) {
		return
	}
	if !g.checkModelAllowed(c, u, br.Model) {
		apiError(c, http.StatusForbidden, "model_not_allowed", "模型未对你开放，请联系管理员")
		return
	}

	// 预扣（reserve）：按 body 估算成本上界计入冻结额；defer 释放（capture），
	// 失败/中断/panic 都会走到，finishUsage 的扣款在冻结释放前完成。
	model, _ := g.modelByName(br.Model)
	g.freezeFor(c, u.ID, model, br, body)
	defer g.unfreeze(c, u.ID)

	cands := g.candidatesFor(c, br.Model)
	if len(cands) == 0 {
		apiError(c, http.StatusServiceUnavailable, "no_channel", "No available upstream channel")
		g.recordError(c, nil, nil, now(), "no_channel")
		return
	}
	sink := &usageSink{}
	c.Set("usage_raw", &sinkHolder{sink: sink})
	g.doProxy(c, cands, convert.ProtoChat, body, br.Stream, sink)
}

// responses POST /v1/responses（OpenAI Responses API）
func (g *Gateway) responses(c *gin.Context) {
	u, ok := g.resolveUser(c)
	if !ok {
		return
	}
	if !g.checkBalance(c, u) {
		return
	}
	br, body, err := parseBody(c)
	if err != nil {
		g.failBodyRead(c, err)
		return
	}
	c.Set("protocol", convert.ProtoResponses)
	c.Set("model_name", br.Model)
	g.recordRawRequest(c, u, body)
	if !g.checkModelExists(c, br.Model) {
		return
	}
	if !g.checkModelAllowed(c, u, br.Model) {
		apiError(c, http.StatusForbidden, "model_not_allowed", "模型未对你开放，请联系管理员")
		return
	}

	// 预扣（reserve）：按 body 估算成本上界计入冻结额；defer 释放（capture），
	// 失败/中断/panic 都会走到，finishUsage 的扣款在冻结释放前完成。
	model, _ := g.modelByName(br.Model)
	g.freezeFor(c, u.ID, model, br, body)
	defer g.unfreeze(c, u.ID)

	cands := g.candidatesFor(c, br.Model)
	if len(cands) == 0 {
		apiError(c, http.StatusServiceUnavailable, "no_channel", "No available upstream channel")
		g.recordError(c, nil, nil, now(), "no_channel")
		return
	}
	sink := &usageSink{}
	c.Set("usage_raw", &sinkHolder{sink: sink})
	g.doProxy(c, cands, convert.ProtoResponses, body, br.Stream, sink)
}

// messages POST /v1/messages（Anthropic Messages API）
func (g *Gateway) messages(c *gin.Context) {
	u, ok := g.resolveUser(c)
	if !ok {
		return
	}
	if !g.checkBalance(c, u) {
		return
	}
	br, body, err := parseBody(c)
	if err != nil {
		g.failBodyRead(c, err)
		return
	}
	c.Set("protocol", convert.ProtoMessages)
	c.Set("model_name", br.Model)
	g.recordRawRequest(c, u, body)
	if !g.checkModelExists(c, br.Model) {
		return
	}
	if !g.checkModelAllowed(c, u, br.Model) {
		apiError(c, http.StatusForbidden, "model_not_allowed", "模型未对你开放，请联系管理员")
		return
	}

	// 预扣（reserve）：按 body 估算成本上界计入冻结额；defer 释放（capture），
	// 失败/中断/panic 都会走到，finishUsage 的扣款在冻结释放前完成。
	model, _ := g.modelByName(br.Model)
	g.freezeFor(c, u.ID, model, br, body)
	defer g.unfreeze(c, u.ID)

	cands := g.candidatesFor(c, br.Model)
	if len(cands) == 0 {
		apiError(c, http.StatusServiceUnavailable, "no_channel", "No available upstream channel")
		g.recordError(c, nil, nil, now(), "no_channel")
		return
	}
	sink := &usageSink{}
	c.Set("usage_raw", &sinkHolder{sink: sink})
	g.doProxy(c, cands, convert.ProtoMessages, body, br.Stream, sink)
}

// usageSinkHolder 桥接：gin context 里保存 sink 引用，供 finishUsage 读取最终 usage。
type sinkHolder struct {
	sink *usageSink
}

// recordRawRequest 记录管理员原始请求体到 context（供 finishUsage 落库）。
// 仅当开关开启且用户为管理员时记录；响应侧以 c.Get("raw_request") 是否非空判断是否需要捕获响应。
func (g *Gateway) recordRawRequest(c *gin.Context, u *store.User, body []byte) {
	if g.logRaw && u.Role == store.RoleAdmin {
		c.Set("raw_request", capRaw(body, rawRequestCap))
	}
}

// apiError 按客户端协议返回错误体（PLANNING §5.1.4）。
func apiError(c *gin.Context, status int, code, message string) {
	if p, _ := c.Get("protocol"); p == convert.ProtoMessages {
		// Anthropic 格式
		c.AbortWithStatusJSON(status, gin.H{
			"type":  "error",
			"error": gin.H{"type": errorTypeFor(status), "message": message},
		})
		return
	}
	// OpenAI 格式
	c.AbortWithStatusJSON(status, gin.H{
		"error": gin.H{
			"message": message,
			"type":    errorTypeFor(status),
			"param":   nil,
			"code":    code,
		},
	})
}

func errorTypeFor(status int) string {
	switch status {
	case http.StatusUnauthorized:
		return "authentication_error"
	case http.StatusForbidden, http.StatusPaymentRequired:
		return "permission_error"
	case http.StatusNotFound, http.StatusBadRequest:
		return "invalid_request_error"
	case http.StatusTooManyRequests:
		return "rate_limit_error"
	default:
		return "api_error"
	}
}