opencatd-open/backend/internal/proxy/gateway.go

package proxy

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"opencatd-open/internal/channel"
	"opencatd-open/internal/dao"
	"opencatd-open/internal/proxy/convert"
	"opencatd-open/internal/store"
	"opencatd-open/internal/usage"
	"opencatd-open/pkg/config"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

type Gateway struct {
	ctx        context.Context
	cfg        *config.Config
	db         *gorm.DB
	wg         *sync.WaitGroup
	httpClient *http.Client

	userDAO    *dao.UserDAO
	apiKeyDAO  *dao.ApiKeyDAO
	usageDAO   *dao.UsageDAO
	dailyDAO   *dao.DailyUsageDAO
	channelSvc *channel.Service
	usageRec   *usage.Recorder
}

func NewGateway(ctx context.Context, cfg *config.Config, db *gorm.DB, wg *sync.WaitGroup, userDAO *dao.UserDAO, apiKeyDAO *dao.ApiKeyDAO, usageDAO *dao.UsageDAO, dailyDAO *dao.DailyUsageDAO) *Gateway {
	client := &http.Client{Timeout: 120 * time.Second}
	if os.Getenv("LOCAL_PROXY") != "" {
		proxyUrl, err := url.Parse(os.Getenv("LOCAL_PROXY"))
		if err == nil {
			tr := &http.Transport{
				Proxy: http.ProxyURL(proxyUrl),
			}
			client.Transport = tr
		}
	}

	return &Gateway{
		ctx:        ctx,
		cfg:        cfg,
		db:         db,
		wg:         wg,
		httpClient: client,
		userDAO:    userDAO,
		apiKeyDAO:  apiKeyDAO,
		usageDAO:   usageDAO,
		dailyDAO:   dailyDAO,
		channelSvc: nil,
	}
}

func (g *Gateway) SetChannelService(svc *channel.Service) {
	g.channelSvc = svc
}

// SetUsageRecorder 注入异步用量记录器；nil 时网关跳过用量上报。
func (g *Gateway) SetUsageRecorder(r *usage.Recorder) {
	g.usageRec = r
}

// Request represents a parsed incoming request
type Request struct {
	Model     string
	Stream    bool
	Protocol  string // "chat", "messages", "responses"
	Body      []byte
	APIKey    *store.APIKey
	UserID    uint64
}

// ParseRequest parses the incoming request and extracts key fields
func (g *Gateway) ParseRequest(c *gin.Context, protocol string) (*Request, error) {
	body, err := io.ReadAll(c.Request.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read body: %w", err)
	}

	apiKey, _ := c.Get("api_key")
	userID, _ := c.Get("user_id")

	req := &Request{
		Protocol: protocol,
		Body:     body,
		UserID:   userID.(uint64),
	}

	if ak, ok := apiKey.(*store.APIKey); ok {
		req.APIKey = ak
	}

	// Parse model and stream based on protocol
	switch protocol {
	case "chat":
		var parsed convert.ChatCompletionRequest
		if err := json.Unmarshal(body, &parsed); err != nil {
			return nil, fmt.Errorf("invalid chat request: %w", err)
		}
		req.Model = parsed.Model
		req.Stream = parsed.Stream
	case "messages":
		var parsed convert.MessagesRequest
		if err := json.Unmarshal(body, &parsed); err != nil {
			return nil, fmt.Errorf("invalid messages request: %w", err)
		}
		req.Model = parsed.Model
		req.Stream = parsed.Stream
	case "responses":
		var parsed convert.ResponsesRequest
		if err := json.Unmarshal(body, &parsed); err != nil {
			return nil, fmt.Errorf("invalid responses request: %w", err)
		}
		req.Model = parsed.Model
		req.Stream = parsed.Stream
	}

	return req, nil
}

// Dispatch routes the request to the appropriate upstream
func (g *Gateway) Dispatch(c *gin.Context, req *Request) {
	if g.channelSvc == nil {
		g.writeError(c, http.StatusBadGateway, "channel service not available")
		return
	}

	ch, err := g.channelSvc.SelectChannel(g.ctx, req.Model)
	if err != nil {
		g.writeError(c, http.StatusBadGateway, err.Error())
		return
	}

	apiKey, err := g.channelSvc.GetAPIKey(ch)
	if err != nil {
		g.writeError(c, http.StatusBadGateway, "failed to decrypt API key")
		return
	}

	// Determine target format: channel declares support for the client protocol
	// then passthrough, otherwise convert to its first supported protocol
	// (chat > messages > responses).
	targetFormat := g.conversionTarget(ch, req.Protocol)
	if targetFormat == "" {
		g.writeError(c, http.StatusBadGateway, fmt.Sprintf("channel %q declares no supported protocol format", ch.Name))
		return
	}

	// Build upstream URL
	upstreamPath := g.getUpstreamPath(targetFormat)
	upstreamURL := ch.UpstreamURL(targetFormat, upstreamPath)

	// Convert request if needed
	var requestBody []byte
	if targetFormat != req.Protocol {
		var err error
		requestBody, err = convert.ConvertRequest(req.Body, req.Protocol, targetFormat)
		if err != nil {
			g.writeError(c, http.StatusBadRequest, "conversion failed: "+err.Error())
			return
		}
	} else {
		requestBody = req.Body
	}

	// Create upstream request
	httpReq, err := http.NewRequestWithContext(g.ctx, "POST", upstreamURL, bytes.NewReader(requestBody))
	if err != nil {
		g.writeError(c, http.StatusBadGateway, "failed to create request")
		return
	}

	// Set headers
	g.setHeaders(httpReq, ch, apiKey, targetFormat)

	// Execute request
	start := time.Now()
	resp, err := g.httpClient.Do(httpReq)
	latency := time.Since(start)
	if err != nil {
		g.channelSvc.RecordFailure(ch.ID)
		g.writeError(c, http.StatusBadGateway, fmt.Sprintf("upstream error: %v (latency: %v)", err, latency))
		return
	}
	defer resp.Body.Close()

	// Record success
	g.channelSvc.RecordSuccess(ch.ID)

	// Handle response
	if resp.StatusCode >= 400 {
		body, _ := io.ReadAll(resp.Body)
		log.Printf("Upstream error: status=%d body=%s", resp.StatusCode, string(body))
		c.Data(resp.StatusCode, "application/json", body)
		return
	}

	// Stream or buffer response
	if req.Stream {
		g.streamResponse(c, resp, req.Protocol, targetFormat)
	} else {
		g.bufferResponse(c, resp, req.Protocol, targetFormat)
	}
}

// conversionTarget 决定客户端协议在渠道上的处理方式：
// 渠道声明支持该协议则直通；否则转为其首选支持协议（chat > messages > responses）。
func (g *Gateway) conversionTarget(ch *store.Channel, clientProto string) string {
	formats := ch.FormatsEffective()
	for _, f := range formats {
		if f == clientProto {
			return clientProto
		}
	}
	for _, p := range []string{convert.ProtoChat, convert.ProtoMessages, convert.ProtoResponses} {
		for _, f := range formats {
			if f == p {
				return p
			}
		}
	}
	return ""
}

func (g *Gateway) getUpstreamPath(protocol string) string {
	switch protocol {
	case "chat":
		return "/chat/completions"
	case "messages":
		return "/messages"
	case "responses":
		return "/responses"
	default:
		return "/chat/completions"
	}
}

func (g *Gateway) setHeaders(req *http.Request, ch *store.Channel, apiKey string, format string) {
	req.Header.Set("Content-Type", "application/json")

	switch ch.Provider {
	case store.ChannelProviderOpenAI, store.ChannelProviderCompatible:
		req.Header.Set("Authorization", "Bearer "+apiKey)
	case store.ChannelProviderAnthropic:
		req.Header.Set("x-api-key", apiKey)
		req.Header.Set("anthropic-version", "2023-06-01")
	}
}


// streamResponse 流式响应：按 \n\n 分块零缓冲转发；跨协议时逐行转换。
func (g *Gateway) streamResponse(c *gin.Context, resp *http.Response, clientProto, upstreamProto string) {
	w := c.Writer
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Status(http.StatusOK)

	flusher, _ := w.(http.Flusher)

	// 跨协议时按行转换；同协议直通（lineConv 为 nil）。
	var lineConv func([]byte) []byte
	if upstreamProto != clientProto {
		lineConv = convert.NewStreamTransformer(upstreamProto, clientProto)
	}

	// 上游原始行按 \n\n 分块，避免把 data 行内的转义换行当成事件边界。
	r := bufio.NewReaderSize(resp.Body, 32*1024)
	for {
		buf := []byte{}
		for {
			line, err := r.ReadSlice('\n')
			if err == bufio.ErrBufferFull {
				buf = append(buf, line...)
				continue
			}
			buf = append(buf, line...)
			if err == io.EOF {
				if len(buf) == 0 {
					return
				}
				if !bytes.HasSuffix(buf, []byte("\n")) {
					buf = append(buf, '\n')
				}
			} else if err != nil {
				log.Printf("stream read error: %v", err)
				return
			}
			if len(buf) >= 2 && bytes.HasSuffix(buf, []byte("\n\n")) {
				break
			}
		}

		out := buf
		if lineConv != nil {
			out = lineConv(buf)
		}
		if len(out) == 0 {
			continue
		}
		if _, err := w.Write(out); err != nil {
			return // 客户端已断开
		}
		if flusher != nil {
			flusher.Flush()
		}
	}
}

func (g *Gateway) bufferResponse(c *gin.Context, resp *http.Response, clientProto, upstreamProto string) {
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		g.writeError(c, http.StatusBadGateway, "failed to read response")
		return
	}

	out := body
	if upstreamProto != clientProto {
		if converted, cerr := convert.ConvertResponse(body, upstreamProto, clientProto); cerr == nil {
			out = converted
		} else {
			// 转换失败时至少剥掉非 JSON 前缀，让客户端能解析出正文
			out = convert.CleanJSON(body)
		}
	} else {
		// 直通：部分上游（如 OpenRouter）的 non-stream 响应在 JSON 前夹带空白/注释
		out = convert.CleanJSON(body)
	}
	c.Data(resp.StatusCode, "application/json", out)
}

func (g *Gateway) writeError(c *gin.Context, status int, message string) {
	protocol := c.GetHeader("X-Protocol")
	if protocol == "" {
		protocol = "chat"
	}

	switch {
	case strings.Contains(c.GetHeader("Accept"), "text/event-stream"):
		c.Header("Content-Type", "text/event-stream")
		c.Status(status)
		fmt.Fprintf(c.Writer, "data: {\"error\":{\"message\":\"%s\"}}\n\n", message)
		fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
	case protocol == "messages":
		c.JSON(status, gin.H{
			"type": "error",
			"error": gin.H{
				"type":    "api_error",
				"message": message,
			},
		})
	default:
		c.JSON(status, gin.H{
			"error": gin.H{
				"message": message,
				"type":    "invalid_request_error",
			},
		})
	}
}