diff --git a/relay/image_handler.go b/relay/image_handler.go index 63ea04d52..99a22a7b6 100644 --- a/relay/image_handler.go +++ b/relay/image_handler.go @@ -6,6 +6,7 @@ import ( "io" "net/http" "strings" + "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" @@ -92,22 +93,37 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type statusCodeMappingStr := c.GetString("status_code_mapping") - // When the client requests streaming, send SSE headers and an initial - // keepalive comment immediately. This prevents reverse-proxies (Nginx - // etc.) from closing the connection due to proxy_read_timeout while the - // backend is waiting for the upstream image API to respond (which can - // take 60+ seconds). The downstream handlers - // (OpenaiImageStreamHandler / OpenaiImageJSONAsStreamHandler) call - // SetEventStreamHeaders again, but that function is idempotent. + // When the client requests streaming, send SSE headers immediately and + // start a periodic keepalive goroutine. Image generation can take 60+ + // seconds; without periodic data the connection will be closed by + // reverse-proxies (Nginx proxy_read_timeout) or the browser. + var keepaliveDone chan struct{} if info.IsStream { helper.SetEventStreamHeaders(c) c.Status(http.StatusOK) - // SSE comment keeps the connection alive through proxies fmt.Fprint(c.Writer, ": keepalive\n\n") _ = helper.FlushWriter(c) + + keepaliveDone = make(chan struct{}) + go func() { + ticker := time.NewTicker(15 * time.Second) + defer ticker.Stop() + for { + select { + case <-ticker.C: + fmt.Fprint(c.Writer, ": keepalive\n\n") + _ = helper.FlushWriter(c) + case <-keepaliveDone: + return + } + } + }() } resp, err := adaptor.DoRequest(c, info, requestBody) + if keepaliveDone != nil { + close(keepaliveDone) + } if err != nil { logImageError(c, info, request, fmt.Sprintf("请求失败: %s", err.Error())) if info.IsStream {