fix(sse): 系统模型 - 恢复 SSE 增量传输

集成 PR #257,保留原 PR 开发者为提交作者。
This commit is contained in:
MasterChiefCN
2026-08-19 20:32:46 +08:00
committed by ddCat
parent d2a35e10d8
commit cd42afdfdd
7 changed files with 233 additions and 14 deletions

View File

@@ -829,13 +829,27 @@ func proxySystemRequest(c *gin.Context, svc *service.Service, user *model.User,
status = model.ApiCallStatusFailed
}
responseLimit := policy.Request.SystemRelayResponseMB << 20
responseBody, readErr := io.ReadAll(io.LimitReader(resp.Body, responseLimit+1))
streamed := isSystemProxyEventStream(resp)
var responseBody []byte
var readErr error
if streamed {
responseBody, readErr = streamSystemProxyResponse(c, resp, responseLimit)
} else {
responseBody, readErr = io.ReadAll(io.LimitReader(resp.Body, responseLimit+1))
}
if readErr != nil {
status = model.ApiCallStatusFailed
errorText = readErr.Error()
_ = svc.MarkBillingUncertain(billingOrderID, "系统渠道响应读取失败,费用状态待核对")
logSystemProxyCall(svc, apiCallLog(user, channel, billingOrderID, capability, protocol, c.Request.Method, path, target, body, c.GetHeader("Content-Type"), status, statusCode, time.Since(startedAt), errorText, concurrencyLimit), nil)
fail(c, http.StatusBadGateway, errors.New("系统渠道响应读取失败"))
billingNote := "系统渠道响应读取失败,费用状态待核对"
if errors.Is(readErr, errSystemProxyResponseTooLarge) {
billingNote = "上游已响应但流式响应体超过限制,费用状态待核对"
}
_ = svc.MarkBillingUncertain(billingOrderID, billingNote)
logSystemProxyCall(svc, apiCallLog(user, channel, billingOrderID, capability, protocol, c.Request.Method, path, target, body, c.GetHeader("Content-Type"), status, statusCode, time.Since(startedAt), errorText, concurrencyLimit), responseBody)
// SSE 响应头已经发送,流中断后只能关闭连接;非流式响应仍返回结构化错误。
if !streamed {
fail(c, http.StatusBadGateway, errors.New("系统渠道响应读取失败"))
}
return
}
if int64(len(responseBody)) > responseLimit {
@@ -855,12 +869,10 @@ func proxySystemRequest(c *gin.Context, svc *service.Service, user *model.User,
} else {
_ = svc.RefundBilling(billingOrderID, "上游明确返回失败")
}
for _, key := range []string{"Content-Type", "Cache-Control", "Content-Disposition"} {
if value := resp.Header.Get(key); value != "" {
c.Header(key, value)
}
if streamed {
return
}
c.Header("X-Content-Type-Options", "nosniff")
copySystemProxyResponseHeaders(c, resp)
c.Data(resp.StatusCode, resp.Header.Get("Content-Type"), responseBody)
}

View File

@@ -0,0 +1,70 @@
package handler
import (
"bytes"
"errors"
"io"
"mime"
"net/http"
"strings"
"github.com/gin-gonic/gin"
)
var errSystemProxyResponseTooLarge = errors.New("系统渠道流式响应超过配置上限")
func isSystemProxyEventStream(resp *http.Response) bool {
mediaType, _, err := mime.ParseMediaType(resp.Header.Get("Content-Type"))
return err == nil && strings.EqualFold(mediaType, "text/event-stream")
}
func copySystemProxyResponseHeaders(c *gin.Context, resp *http.Response) {
for _, key := range []string{"Content-Type", "Cache-Control", "Content-Disposition"} {
if value := resp.Header.Get(key); value != "" {
c.Header(key, value)
}
}
c.Header("X-Content-Type-Options", "nosniff")
}
func streamSystemProxyResponse(c *gin.Context, resp *http.Response, responseLimit int64) ([]byte, error) {
copySystemProxyResponseHeaders(c, resp)
// 同时通知镜像内和可能存在的外层 Nginx,不得重新缓冲模型事件流。
c.Header("Cache-Control", "no-cache, no-store, no-transform")
c.Header("X-Accel-Buffering", "no")
c.Status(resp.StatusCode)
c.Writer.WriteHeaderNow()
c.Writer.Flush()
var captured bytes.Buffer
buffer := make([]byte, 32<<10)
remaining := responseLimit
for {
readLimit := int64(len(buffer))
if remaining < readLimit {
readLimit = remaining + 1
}
if readLimit < 1 {
readLimit = 1
}
read, readErr := resp.Body.Read(buffer[:int(readLimit)])
if read > 0 {
if int64(read) > remaining {
return captured.Bytes(), errSystemProxyResponseTooLarge
}
chunk := buffer[:read]
_, _ = captured.Write(chunk)
if _, writeErr := c.Writer.Write(chunk); writeErr != nil {
return captured.Bytes(), writeErr
}
c.Writer.Flush()
remaining -= int64(read)
}
if readErr != nil {
if errors.Is(readErr, io.EOF) {
return captured.Bytes(), nil
}
return captured.Bytes(), readErr
}
}
}

View File

@@ -0,0 +1,85 @@
package handler
import (
"bufio"
"errors"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/gin-gonic/gin"
)
func TestSystemProxyEventStreamFlushesBeforeUpstreamCompletes(t *testing.T) {
gin.SetMode(gin.TestMode)
release := make(chan struct{})
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
_, _ = io.WriteString(w, "data: first\n\n")
w.(http.Flusher).Flush()
<-release
_, _ = io.WriteString(w, "data: second\n\n")
}))
defer upstream.Close()
proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
resp, err := http.Get(upstream.URL)
if err != nil {
http.Error(w, err.Error(), http.StatusBadGateway)
return
}
defer resp.Body.Close()
context, _ := gin.CreateTestContext(w)
context.Request = r
_, _ = streamSystemProxyResponse(context, resp, 1<<20)
}))
defer proxy.Close()
client := &http.Client{Timeout: 3 * time.Second}
response, err := client.Get(proxy.URL)
if err != nil {
close(release)
t.Fatal(err)
}
if response.Header.Get("X-Accel-Buffering") != "no" {
close(release)
_ = response.Body.Close()
t.Fatalf("X-Accel-Buffering = %q", response.Header.Get("X-Accel-Buffering"))
}
line, err := bufio.NewReader(response.Body).ReadString('\n')
if err != nil {
close(release)
_ = response.Body.Close()
t.Fatal(err)
}
if line != "data: first\n" {
close(release)
_ = response.Body.Close()
t.Fatalf("first streamed line = %q", line)
}
close(release)
_ = response.Body.Close()
}
func TestSystemProxyEventStreamEnforcesResponseLimit(t *testing.T) {
gin.SetMode(gin.TestMode)
response := httptest.NewRecorder()
context, _ := gin.CreateTestContext(response)
context.Request = httptest.NewRequest(http.MethodPost, "/api/ai/system/channel/responses", nil)
upstream := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader("data: too large\n\n")),
}
captured, err := streamSystemProxyResponse(context, upstream, 4)
if !errors.Is(err, errSystemProxyResponseTooLarge) {
t.Fatalf("error = %v", err)
}
if len(captured) != 0 || response.Body.Len() != 0 {
t.Fatalf("captured = %q, response = %q", captured, response.Body.String())
}
}