AI 网关切换 OpenAI 兼容面并移除 chat 端点,新增模型黑白名单

internal/api/aigateway.go

@@ -1,13 +1,15 @@
 package api
 
 import (
+	"bufio"
+	"bytes"
 	"crypto/rand"
 	"encoding/hex"
 	"encoding/json"
 	"errors"
-	"fmt"
 	"io"
 	"net/http"
+	"slices"
 	"strings"
 	"time"
 
@@ -86,6 +88,25 @@ func keyGroup(c *gin.Context) string {
 	return ""
 }
 
+// keyModels 取当前请求密钥的模型白名单(空 = 不限模型)。
+func keyModels(c *gin.Context) []string {
+	if key, ok := c.Get(aiKeyCtx); ok {
+		return key.(*model.AiKey).Models
+	}
+	return nil
+}
+
+// checkKeyModel 校验请求模型是否在密钥白名单内;拒绝时按端点协议返回
+// 404 model_not_found(与未知模型同口径,不泄露密钥配置细节)并返回 false。
+func checkKeyModel(c *gin.Context, modelName string) bool {
+	allowed := keyModels(c)
+	if len(allowed) == 0 || slices.Contains(allowed, modelName) {
+		return true
+	}
+	aiError(c, http.StatusNotFound, "model_not_found", "模型不存在或无权访问: "+modelName)
+	return false
+}
+
 // logEntry 组装调用日志骨架;usage / status 由调用方补齐。
 func (h *aiGatewayHandler) logEntry(c *gin.Context, endpoint, modelName string, stream bool, meta service.ChatMeta, start time.Time) model.AiCallLog {
 	entry := model.AiCallLog{
@@ -102,21 +123,6 @@ func (h *aiGatewayHandler) logEntry(c *gin.Context, endpoint, modelName string,
 }
 
 // validateIR 做协议无关的入参检查(模型名、消息、不支持的内容块)。
-func validateIR(ir aiwire.ChatRequest) error {
-	if strings.TrimSpace(ir.Model) == "" {
-		return fmt.Errorf("model 不能为空")
-	}
-	if len(ir.Messages) == 0 {
-		return fmt.Errorf("messages 不能为空")
-	}
-	for _, m := range ir.Messages {
-		if m.Content.HasUnsupported() {
-			return service.ErrAiUnsupportedBlock
-		}
-	}
-	return nil
-}
-
 // maybeLogContent 在调用密钥显式开启内容日志且未过期时写正文(红线例外);
 // respBody 为 nil 时只记请求(流式响应与向量结果不记录);callLogID 关联同次调用日志。
 func (h *aiGatewayHandler) maybeLogContent(c *gin.Context, callLogID uint, endpoint, modelName string, stream bool, reqBody, respBody any) {
@@ -147,47 +153,6 @@ func (h *aiGatewayHandler) logFailure(c *gin.Context, entry model.AiCallLog, req
 	h.maybeLogContent(c, callID, entry.Endpoint, entry.Model, entry.Stream, reqBody, nil)
 }
 
-// chatCompletions 是 OpenAI /ai/v1/chat/completions 端点。
-//
-// @Summary  OpenAI 兼容对话补全
-// @Tags     AI 网关
-// @Param    body body object true "OpenAI chat/completions 请求体(支持 stream)"
-// @Success  200 {object} map[string]any "OpenAI 兼容响应(流式为 SSE)"
-// @Router   /ai/v1/chat/completions [post]
-func (h *aiGatewayHandler) chatCompletions(c *gin.Context) {
-	var ir aiwire.ChatRequest
-	if err := c.ShouldBindJSON(&ir); err != nil {
-		aiError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
-		return
-	}
-	if err := validateIR(ir); err != nil {
-		aiError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
-		return
-	}
-	if ir.Stream {
-		h.streamOpenAI(c, ir)
-		return
-	}
-	start := time.Now()
-	resp, meta, err := h.gw.Chat(c.Request.Context(), ir, keyGroup(c))
-	entry := h.logEntry(c, "openai", ir.Model, false, meta, start)
-	if err != nil {
-		upstreamError(c, err)
-		entry.ErrMsg = err.Error()
-		h.logFailure(c, entry, ir)
-		return
-	}
-	resp.ID = aiRandID("chatcmpl-")
-	if resp.Created == 0 {
-		resp.Created = time.Now().Unix()
-	}
-	entry.Status = http.StatusOK
-	fillUsage(&entry, resp.Usage)
-	callID := h.gw.LogCall(entry)
-	h.maybeLogContent(c, callID, "openai", ir.Model, false, ir, resp)
-	c.JSON(http.StatusOK, resp)
-}
-
 func fillUsage(entry *model.AiCallLog, u *aiwire.Usage) {
 	if u == nil {
 		return
@@ -217,6 +182,9 @@ func (h *aiGatewayHandler) embeddings(c *gin.Context) {
 		aiError(c, http.StatusBadRequest, "invalid_request_error", "仅支持 encoding_format=float")
 		return
 	}
+	if !checkKeyModel(c, req.Model) {
+		return
+	}
 	start := time.Now()
 	resp, meta, err := h.gw.Embeddings(c.Request.Context(), req, keyGroup(c))
 	entry := h.logEntry(c, "embeddings", req.Model, false, meta, start)
@@ -235,44 +203,6 @@ func (h *aiGatewayHandler) embeddings(c *gin.Context) {
 	c.JSON(http.StatusOK, resp)
 }
 
-// streamOpenAI 以 OpenAI SSE 直通流式响应,末尾发 [DONE]。
-func (h *aiGatewayHandler) streamOpenAI(c *gin.Context, ir aiwire.ChatRequest) {
-	start := time.Now()
-	stream, meta, err := h.gw.OpenStream(c.Request.Context(), ir, keyGroup(c))
-	entry := h.logEntry(c, "openai", ir.Model, true, meta, start)
-	if err != nil {
-		upstreamError(c, err)
-		entry.ErrMsg = err.Error()
-		h.logFailure(c, entry, ir)
-		return
-	}
-	defer stream.Close()
-	sseHeaders(c)
-	id, created := aiRandID("chatcmpl-"), time.Now().Unix()
-	var usage *aiwire.Usage
-	for {
-		chunk, err := stream.Next()
-		if err != nil {
-			if !errors.Is(err, io.EOF) {
-				entry.ErrMsg = err.Error()
-			}
-			break
-		}
-		chunk.ID, chunk.Created = id, created
-		if chunk.Usage != nil {
-			usage = chunk.Usage
-		}
-		writeSSEData(c, chunk)
-	}
-	c.Writer.WriteString("data: [DONE]\n\n")
-	c.Writer.Flush()
-	entry.Status = http.StatusOK
-	entry.LatencyMs = time.Since(start).Milliseconds()
-	fillUsage(&entry, usage)
-	callID := h.gw.LogCall(entry)
-	h.maybeLogContent(c, callID, "openai", ir.Model, true, ir, nil)
-}
-
 func sseHeaders(c *gin.Context) {
 	c.Header("Content-Type", "text/event-stream")
 	c.Header("Cache-Control", "no-cache")
@@ -280,22 +210,11 @@ func sseHeaders(c *gin.Context) {
 	c.Writer.Flush()
 }
 
-func writeSSEData(c *gin.Context, v any) {
-	b, err := json.Marshal(v)
-	if err != nil {
-		return
-	}
-	c.Writer.WriteString("data: ")
-	c.Writer.Write(b)
-	c.Writer.WriteString("\n\n")
-	c.Writer.Flush()
-}
-
 // messages 是 Anthropic /ai/v1/messages 端点。
 //
 // @Summary  Anthropic Messages 兼容端点
 // @Tags     AI 网关
-// @Param    body body object true "Anthropic messages 请求体(支持 stream)"
+// @Param    body body object true "Anthropic messages 请求体(支持 stream;经 OCI OpenAI 兼容面直通)"
 // @Success  200 {object} map[string]any "Anthropic 兼容响应(流式为 SSE)"
 // @Router   /ai/v1/messages [post]
 func (h *aiGatewayHandler) messages(c *gin.Context) {
@@ -308,69 +227,92 @@ func (h *aiGatewayHandler) messages(c *gin.Context) {
 		aiError(c, http.StatusBadRequest, "invalid_request_error", "max_tokens 必填且需大于 0")
 		return
 	}
-	ir, err := service.AnthropicToIR(req)
+	if len(req.Messages) == 0 {
+		aiError(c, http.StatusBadRequest, "invalid_request_error", "messages 不能为空")
+		return
+	}
+	if !checkKeyModel(c, req.Model) {
+		return
+	}
+	body, err := service.AnthropicToResponsesBody(req)
 	if err != nil {
 		aiError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
 		return
 	}
-	if err := validateIR(ir); err != nil {
-		aiError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
-		return
-	}
 	if req.Stream {
-		h.streamAnthropic(c, ir)
+		h.streamAnthropic(c, body, req)
 		return
 	}
 	start := time.Now()
-	resp, meta, err := h.gw.Chat(c.Request.Context(), ir, keyGroup(c))
-	entry := h.logEntry(c, "anthropic", ir.Model, false, meta, start)
+	payload, meta, err := h.gw.RespPassthrough(c.Request.Context(), body, req.Model, keyGroup(c))
+	entry := h.logEntry(c, "anthropic", req.Model, false, meta, start)
 	if err != nil {
 		upstreamError(c, err)
 		entry.ErrMsg = err.Error()
 		h.logFailure(c, entry, req)
 		return
 	}
+	out, err := service.ResponsesToAnthropic(payload, aiRandID("msg_"))
+	if err != nil {
+		aiError(c, http.StatusBadGateway, "api_error", err.Error())
+		entry.ErrMsg = err.Error()
+		h.logFailure(c, entry, req)
+		return
+	}
 	entry.Status = http.StatusOK
-	fillUsage(&entry, resp.Usage)
+	fillUsage(&entry, service.RespPassthroughUsage(payload))
 	callID := h.gw.LogCall(entry)
-	out := service.IRRespToAnthropic(resp, aiRandID("msg_"))
-	h.maybeLogContent(c, callID, "anthropic", ir.Model, false, req, out)
+	h.maybeLogContent(c, callID, "anthropic", req.Model, false, req, out)
 	c.JSON(http.StatusOK, out)
 }
 
-// streamAnthropic 经状态机把 IR chunk 流聚合为 Anthropic SSE 事件流。
-func (h *aiGatewayHandler) streamAnthropic(c *gin.Context, ir aiwire.ChatRequest) {
+// streamAnthropic 流式直通并桥接:上游 Responses SSE 事件经 AnthRespBridge
+// 转为 Anthropic 事件序列逐块写出;usage 从 completed 事件记账。
+func (h *aiGatewayHandler) streamAnthropic(c *gin.Context, body []byte, req aiwire.MessagesRequest) {
 	start := time.Now()
-	stream, meta, err := h.gw.OpenStream(c.Request.Context(), ir, keyGroup(c))
-	entry := h.logEntry(c, "anthropic", ir.Model, true, meta, start)
+	upstream, meta, err := h.gw.RespPassthroughStream(c.Request.Context(), body, req.Model, keyGroup(c))
+	entry := h.logEntry(c, "anthropic", req.Model, true, meta, start)
 	if err != nil {
 		upstreamError(c, err)
 		entry.ErrMsg = err.Error()
-		h.logFailure(c, entry, ir)
+		h.logFailure(c, entry, req)
 		return
 	}
-	defer stream.Close()
+	defer upstream.Close()
 	sseHeaders(c)
-	st := service.NewAnthStream(aiRandID("msg_"), ir.Model)
-	for {
-		chunk, err := stream.Next()
-		if err != nil {
-			if !errors.Is(err, io.EOF) {
-				entry.ErrMsg = err.Error()
-			}
-			break
-		}
-		writeAnthEvents(c, st.Feed(chunk))
+	bridge := service.NewAnthRespBridge(aiRandID("msg_"), req.Model)
+	if err := forwardAnthSSE(c, upstream, bridge); err != nil {
+		entry.ErrMsg = err.Error()
 	}
-	writeAnthEvents(c, st.Finish())
+	writeAnthEvents(c, bridge.Finish())
 	entry.Status = http.StatusOK
 	entry.LatencyMs = time.Since(start).Milliseconds()
-	usage := st.Usage()
+	usage := bridge.Usage()
 	entry.PromptTokens, entry.CompletionTokens = usage.InputTokens, usage.OutputTokens
 	entry.TotalTokens = usage.InputTokens + usage.OutputTokens
 	entry.CachedTokens = usage.CacheReadInputTokens
 	callID := h.gw.LogCall(entry)
-	h.maybeLogContent(c, callID, "anthropic", ir.Model, true, ir, nil)
+	h.maybeLogContent(c, callID, "anthropic", req.Model, true, req, nil)
+}
+
+// forwardAnthSSE 逐行读上游 SSE,把 data 行喂给桥并即时写出转换后的事件。
+func forwardAnthSSE(c *gin.Context, upstream io.Reader, bridge *service.AnthRespBridge) error {
+	reader := bufio.NewReader(upstream)
+	for {
+		line, err := reader.ReadBytes('\n')
+		if len(line) > 0 {
+			trimmed := bytes.TrimSpace(line)
+			if data, ok := bytes.CutPrefix(trimmed, []byte("data: ")); ok {
+				writeAnthEvents(c, bridge.Feed(data))
+			}
+		}
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				return nil
+			}
+			return err
+		}
+	}
 }
 
 func writeAnthEvents(c *gin.Context, events []service.AnthEvent) {
@@ -398,6 +340,15 @@ func (h *aiGatewayHandler) listModels(c *gin.Context) {
 		aiError(c, http.StatusInternalServerError, "api_error", "查询模型列表失败")
 		return
 	}
+	if allowed := keyModels(c); len(allowed) > 0 {
+		kept := make([]aiwire.Model, 0, len(list.Data))
+		for _, m := range list.Data {
+			if slices.Contains(allowed, m.ID) {
+				kept = append(kept, m)
+			}
+		}
+		list.Data = kept
+	}
 	c.JSON(http.StatusOK, list)
 }
 
@@ -405,31 +356,45 @@ func (h *aiGatewayHandler) listModels(c *gin.Context) {
 //
 // @Summary  OpenAI Responses 兼容端点
 // @Tags     AI 网关
-// @Param    body body object true "OpenAI responses 请求体(支持 stream)"
+// @Param    body body object true "OpenAI responses 请求体(支持 stream;web_search/x_search 服务端工具仅非流式)"
 // @Success  200 {object} map[string]any "OpenAI 兼容响应(流式为 SSE)"
 // @Router   /ai/v1/responses [post]
 func (h *aiGatewayHandler) responses(c *gin.Context) {
-	var req aiwire.RespRequest
-	if err := c.ShouldBindJSON(&req); err != nil {
-		aiError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
-		return
-	}
-	ir, err := service.ResponsesToIR(req)
+	raw, err := c.GetRawData()
 	if err != nil {
 		aiError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
 		return
 	}
-	if err := validateIR(ir); err != nil {
+	var req aiwire.RespRequest
+	if err := json.Unmarshal(raw, &req); err != nil {
+		aiError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
+		return
+	}
+	if !checkKeyModel(c, req.Model) {
+		return
+	}
+	h.responsesPassthrough(c, raw, req)
+}
+
+// responsesPassthrough 直通链路(唯一上游):原始 body 改写后直达
+// OCI /actions/v1/responses,响应原样透传。
+func (h *aiGatewayHandler) responsesPassthrough(c *gin.Context, raw []byte, req aiwire.RespRequest) {
+	if err := service.RespPassthroughValidate(req); err != nil {
+		aiError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
+		return
+	}
+	body, err := service.RespPassthroughBody(raw)
+	if err != nil {
 		aiError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
 		return
 	}
 	if req.Stream {
-		h.streamResponses(c, ir)
+		h.responsesPassthroughStream(c, body, req)
 		return
 	}
 	start := time.Now()
-	resp, meta, err := h.gw.Chat(c.Request.Context(), ir, keyGroup(c))
-	entry := h.logEntry(c, "responses", ir.Model, false, meta, start)
+	payload, meta, err := h.gw.RespPassthrough(c.Request.Context(), body, req.Model, keyGroup(c))
+	entry := h.logEntry(c, "responses", req.Model, false, meta, start)
 	if err != nil {
 		upstreamError(c, err)
 		entry.ErrMsg = err.Error()
@@ -437,54 +402,61 @@ func (h *aiGatewayHandler) responses(c *gin.Context) {
 		return
 	}
 	entry.Status = http.StatusOK
-	fillUsage(&entry, resp.Usage)
+	fillUsage(&entry, service.RespPassthroughUsage(payload))
 	callID := h.gw.LogCall(entry)
-	out := service.IRRespToResponses(resp, aiRandID("resp_"), time.Now().Unix())
-	h.maybeLogContent(c, callID, "responses", ir.Model, false, req, out)
-	c.JSON(http.StatusOK, out)
+	h.maybeLogContent(c, callID, "responses", req.Model, false, req, json.RawMessage(payload))
+	c.Data(http.StatusOK, "application/json; charset=utf-8", payload)
 }
 
-// streamResponses 经状态机把 IR chunk 流聚合为 Responses 语义事件流。
-func (h *aiGatewayHandler) streamResponses(c *gin.Context, ir aiwire.ChatRequest) {
+// responsesPassthroughStream 流式直通:SSE 事件原样转发(推理增量等直达客户端),
+// 逐行扫描 completed 事件提取 usage 记账。
+func (h *aiGatewayHandler) responsesPassthroughStream(c *gin.Context, body []byte, req aiwire.RespRequest) {
 	start := time.Now()
-	stream, meta, err := h.gw.OpenStream(c.Request.Context(), ir, keyGroup(c))
-	entry := h.logEntry(c, "responses", ir.Model, true, meta, start)
+	upstream, meta, err := h.gw.RespPassthroughStream(c.Request.Context(), body, req.Model, keyGroup(c))
+	entry := h.logEntry(c, "responses", req.Model, true, meta, start)
 	if err != nil {
 		upstreamError(c, err)
 		entry.ErrMsg = err.Error()
-		h.logFailure(c, entry, ir)
+		h.logFailure(c, entry, req)
 		return
 	}
-	defer stream.Close()
+	defer upstream.Close()
 	sseHeaders(c)
-	st := service.NewRespStream(aiRandID("resp_"), ir.Model, time.Now().Unix())
-	for {
-		chunk, err := stream.Next()
-		if err != nil {
-			if !errors.Is(err, io.EOF) {
-				entry.ErrMsg = err.Error()
-			}
-			break
-		}
-		writeRespEvents(c, st.Feed(chunk))
+	usage, err := forwardSSE(c, upstream)
+	if err != nil {
+		entry.ErrMsg = err.Error()
 	}
-	writeRespEvents(c, st.Finish())
 	entry.Status = http.StatusOK
 	entry.LatencyMs = time.Since(start).Milliseconds()
-	fillUsage(&entry, st.Usage())
+	fillUsage(&entry, usage)
 	callID := h.gw.LogCall(entry)
-	h.maybeLogContent(c, callID, "responses", ir.Model, true, ir, nil)
+	h.maybeLogContent(c, callID, "responses", req.Model, true, req, nil)
 }
 
-func writeRespEvents(c *gin.Context, events []service.RespEvent) {
-	for _, ev := range events {
-		b, err := json.Marshal(ev.Data)
+// forwardSSE 把上游 SSE 逐行转发给客户端,空行(事件边界)即 flush;
+// 顺带从 data 行提取 response.completed 的 usage。
+func forwardSSE(c *gin.Context, upstream io.Reader) (*aiwire.Usage, error) {
+	reader := bufio.NewReader(upstream)
+	var usage *aiwire.Usage
+	for {
+		line, err := reader.ReadBytes('\n')
+		if len(line) > 0 {
+			c.Writer.Write(line)
+			trimmed := bytes.TrimSpace(line)
+			if len(trimmed) == 0 {
+				c.Writer.Flush()
+			} else if data, ok := bytes.CutPrefix(trimmed, []byte("data: ")); ok {
+				if u := service.RespStreamCompletedUsage(data); u != nil {
+					usage = u
+				}
+			}
+		}
 		if err != nil {
-			continue
+			c.Writer.Flush()
+			if errors.Is(err, io.EOF) {
+				return usage, nil
+			}
+			return usage, err
 		}
-		c.Writer.WriteString("event: " + ev.Event + "\ndata: ")
-		c.Writer.Write(b)
-		c.Writer.WriteString("\n\n")
 	}
-	c.Writer.Flush()
 }