@@ -41,6 +41,11 @@ type Gateway struct {
modelDAO * dao . ModelDAO
channelSvc * channel . Service
usageRec * usage . Recorder
// 原始请求/响应记录开关(系统配置 log_raw_requests,带 TTL 缓存避免每次查库)。
rawLogMu sync . Mutex
rawLogVal bool
rawLogSet time . Time
}
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 {
@@ -79,6 +84,22 @@ func (g *Gateway) SetUsageRecorder(r *usage.Recorder) {
g . usageRec = r
}
// rawLogEnabled 读取系统配置 log_raw_requests( 10s TTL 缓存),决定是否记录原始请求/响应。
func ( g * Gateway ) rawLogEnabled ( ) bool {
g . rawLogMu . Lock ( )
defer g . rawLogMu . Unlock ( )
if time . Since ( g . rawLogSet ) < 10 * time . Second {
return g . rawLogVal
}
var sc store . SystemConfig
g . rawLogVal = false
if err := g . db . Where ( "key = ?" , "log_raw_requests" ) . First ( & sc ) . Error ; err == nil {
g . rawLogVal = strings . TrimSpace ( sc . Value ) == "true"
}
g . rawLogSet = time . Now ( )
return g . rawLogVal
}
// generateRequestID 生成请求级唯一 ID,用于用量明细关联与排障。
func generateRequestID ( ) string {
b := make ( [ ] byte , 12 )
@@ -98,6 +119,10 @@ type Request struct {
UserID uint64
KeyID uint64
RequestID string
CaptureRaw bool // 原始请求/响应记录(管理员 + 系统开关开启)
rawBuf * strings . Builder // 上游原始响应累积器(仅 CaptureRaw 时非 nil)
}
// ParseRequest parses the incoming request and extracts key fields
@@ -109,6 +134,7 @@ func (g *Gateway) ParseRequest(c *gin.Context, protocol string) (*Request, error
apiKey , _ := c . Get ( "api_key" )
userID , _ := c . Get ( "user_id" )
userRole , _ := c . Get ( "user_role" )
req := & Request {
Protocol : protocol ,
@@ -124,6 +150,11 @@ func (g *Gateway) ParseRequest(c *gin.Context, protocol string) (*Request, error
req . APIKey = ak
}
// 原始请求/响应记录:仅管理员 且 系统开关 log_raw_requests 开启。
if role , _ := userRole . ( string ) ; role == store . RoleAdmin && g . rawLogEnabled ( ) {
req . CaptureRaw = true
}
// Parse model and stream based on protocol
switch protocol {
case "chat" :
@@ -160,6 +191,12 @@ func (g *Gateway) Dispatch(c *gin.Context, req *Request) {
return
}
// 原始请求/响应捕获:仅管理员 + 系统开关开启(req.CaptureRaw 已在 ParseRequest 判定)。
// 客户端原始请求体即 req.Body;上游原始响应由 stream/bufferResponse 累积进 rawBuf。
if req . CaptureRaw {
req . rawBuf = & strings . Builder { }
}
cands := g . channelSvc . Candidates ( req . Model )
// 内存健康过滤:连续失败进入 cooldown 的渠道不再尝试(渠道级健康自愈靠冷却过期)。
cands = g . channelSvc . FilterHealthy ( cands )
@@ -245,6 +282,9 @@ func (g *Gateway) Dispatch(c *gin.Context, req *Request) {
body , _ := io . ReadAll ( resp . Body )
resp . Body . Close ( )
log . Printf ( "Upstream error: status=%d body=%s" , resp . StatusCode , string ( body ) )
if req . rawBuf != nil {
req . rawBuf . Write ( body )
}
g . recordUsage ( req , cand , ch , usage . Event {
IsError : true ,
ErrorCode : fmt . Sprintf ( "upstream_%d" , resp . StatusCode ) ,
@@ -264,9 +304,9 @@ func (g *Gateway) Dispatch(c *gin.Context, req *Request) {
// Stream or buffer response; tok 从上游响应(SSE usage 块或非流式 JSON)提取。
var tok convert . TokenUsage
if req . Stream {
tok = g . streamResponse ( c , resp , req . Protocol , targetFormat )
tok = g . streamResponse ( c , resp , req . Protocol , targetFormat , req . rawBuf )
} else {
tok = g . bufferResponse ( c , resp , req . Protocol , targetFormat )
tok = g . bufferResponse ( c , resp , req . Protocol , targetFormat , req . rawBuf )
}
resp . Body . Close ( )
// 成功记录:用量 + 定价计费。
@@ -320,6 +360,13 @@ func (g *Gateway) recordUsage(req *Request, cand *channel.Candidate, ch *store.C
ev . CompletionTokens = tok . OutputTokens
ev . CacheReadTokens = tok . CacheReadTokens
ev . CacheCreationTokens = tok . CacheCreationTokens
// 原始请求/响应(仅管理员+开关开启时捕获)。
if req . CaptureRaw {
ev . RawRequest = string ( req . Body )
if req . rawBuf != nil {
ev . RawResponse = req . rawBuf . String ( )
}
}
// 定价与成本(价格按每百万 token 的 USD 单价)。
// 成本口径:非缓存输入 × 输入价 + 缓存读 × 缓存价 + 缓存写与输出 × 输出价。
if ev . ModelID != 0 {
@@ -385,7 +432,8 @@ func (g *Gateway) setHeaders(req *http.Request, ch *store.Channel, apiKey string
// streamResponse 流式响应:按 \n\n 分块零缓冲转发;跨协议时逐行转换。
// 返回从上游 SSE usage 块累计的 token 用量(按上游协议解析)。
func ( g * Gateway ) streamResponse ( c * gin . Context , resp * http . Response , clientProto , upstreamProto string ) convert . TokenUsage {
// capture 非 nil 时把上游原始行累积进去(原始响应记录)。
func ( g * Gateway ) streamResponse ( c * gin . Context , resp * http . Response , clientProto , upstreamProto string , capture * strings . Builder ) convert . TokenUsage {
w := c . Writer
c . Header ( "Content-Type" , "text/event-stream" )
c . Header ( "Cache-Control" , "no-cache" )
@@ -429,6 +477,11 @@ func (g *Gateway) streamResponse(c *gin.Context, resp *http.Response, clientProt
}
}
// 原始响应捕获(仅管理员+开关开启时启用)。
if capture != nil {
capture . Write ( buf )
}
// 先解析用量(data: {...} 行),再决定转发内容。
for _ , data := range sseDataPayloads ( buf ) {
accum . Feed ( data , upstreamProto )
@@ -493,13 +546,18 @@ func sseDataPayloads(chunk []byte) [][]byte {
return out
}
func ( g * Gateway ) bufferResponse ( c * gin . Context , resp * http . Response , clientProto , upstreamProto string ) convert . TokenUsage {
func ( g * Gateway ) bufferResponse ( c * gin . Context , resp * http . Response , clientProto , upstreamProto string , capture * strings . Builder ) convert . TokenUsage {
body , err := io . ReadAll ( resp . Body )
if err != nil {
g . writeError ( c , http . StatusBadGateway , "failed to read response" )
return convert . TokenUsage { }
}
// 原始响应捕获(仅管理员+开关开启时启用)。
if capture != nil {
capture . Write ( body )
}
// 用量从上游原始响应体提取(先于转换,转换会改字段名)。
tok , _ := convert . ExtractUsageJSON ( body , upstreamProto )