mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-21 14:19:18 +08:00
Merge pull request #4307 from wp-a/fix/openai-first-output-timeout
fix(openai): bound native responses first output wait
This commit is contained in:
@@ -748,6 +748,11 @@ type GatewayConfig struct {
|
||||
// OpenAIResponseHeaderTimeout: OpenAI/Codex 上游等待响应头的超时时间(秒),0表示无超时
|
||||
// OpenAI/Codex 请求可能在上游排队较久;默认不使用通用响应头超时截断。
|
||||
OpenAIResponseHeaderTimeout int `mapstructure:"openai_response_header_timeout"`
|
||||
// OpenAIFirstOutputTimeoutSeconds: native HTTP Responses 首个语义输出超时(秒),0表示禁用。
|
||||
OpenAIFirstOutputTimeoutSeconds int `mapstructure:"openai_first_output_timeout_seconds"`
|
||||
// OpenAIHighEffortFirstOutputTimeoutSeconds: high/xhigh/max 推理的首个语义输出超时(秒)。
|
||||
// 0 表示回退到 OpenAIFirstOutputTimeoutSeconds。
|
||||
OpenAIHighEffortFirstOutputTimeoutSeconds int `mapstructure:"openai_high_effort_first_output_timeout_seconds"`
|
||||
// 请求体最大字节数,用于网关请求体大小限制
|
||||
MaxBodySize int64 `mapstructure:"max_body_size"`
|
||||
// 非流式上游响应体读取上限(字节),用于防止无界读取导致内存放大
|
||||
@@ -1944,6 +1949,8 @@ func setDefaults() {
|
||||
// Gateway
|
||||
viper.SetDefault("gateway.response_header_timeout", 600) // 600秒(10分钟)等待上游响应头,LLM高负载时可能排队较久
|
||||
viper.SetDefault("gateway.openai_response_header_timeout", 0)
|
||||
viper.SetDefault("gateway.openai_first_output_timeout_seconds", 0)
|
||||
viper.SetDefault("gateway.openai_high_effort_first_output_timeout_seconds", 0)
|
||||
viper.SetDefault("gateway.log_upstream_error_body", true)
|
||||
viper.SetDefault("gateway.log_upstream_error_body_max_bytes", 2048)
|
||||
viper.SetDefault("gateway.inject_beta_for_apikey", false)
|
||||
@@ -2651,6 +2658,14 @@ func (c *Config) Validate() error {
|
||||
if c.Gateway.OpenAIResponseHeaderTimeout < 0 {
|
||||
return fmt.Errorf("gateway.openai_response_header_timeout must be non-negative")
|
||||
}
|
||||
if c.Gateway.OpenAIFirstOutputTimeoutSeconds < 0 || c.Gateway.OpenAIFirstOutputTimeoutSeconds > 600 ||
|
||||
(c.Gateway.OpenAIFirstOutputTimeoutSeconds > 0 && c.Gateway.OpenAIFirstOutputTimeoutSeconds < 30) {
|
||||
return fmt.Errorf("gateway.openai_first_output_timeout_seconds must be 0 or between 30-600 seconds")
|
||||
}
|
||||
if c.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds < 0 || c.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds > 1800 ||
|
||||
(c.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds > 0 && c.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds < 30) {
|
||||
return fmt.Errorf("gateway.openai_high_effort_first_output_timeout_seconds must be 0 or between 30-1800 seconds")
|
||||
}
|
||||
if strings.TrimSpace(c.Gateway.ConnectionPoolIsolation) != "" {
|
||||
switch c.Gateway.ConnectionPoolIsolation {
|
||||
case ConnectionPoolIsolationProxy, ConnectionPoolIsolationAccount, ConnectionPoolIsolationAccountProxy:
|
||||
|
||||
@@ -98,6 +98,34 @@ func TestLoadDefaultSchedulingConfig(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadDefaultOpenAIFirstOutputTimeoutsDisabled(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, cfg.Gateway.OpenAIFirstOutputTimeoutSeconds)
|
||||
require.Zero(t, cfg.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds)
|
||||
}
|
||||
|
||||
func TestLoadOpenAIFirstOutputTimeoutsFromEnv(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("GATEWAY_OPENAI_FIRST_OUTPUT_TIMEOUT_SECONDS", "90")
|
||||
t.Setenv("GATEWAY_OPENAI_HIGH_EFFORT_FIRST_OUTPUT_TIMEOUT_SECONDS", "240")
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 90, cfg.Gateway.OpenAIFirstOutputTimeoutSeconds)
|
||||
require.Equal(t, 240, cfg.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds)
|
||||
}
|
||||
|
||||
func TestValidateOpenAIFirstOutputTimeoutMinimum(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
cfg.Gateway.OpenAIFirstOutputTimeoutSeconds = 30
|
||||
require.NoError(t, cfg.Validate())
|
||||
}
|
||||
|
||||
func TestLoadDefaultOpenAIWSConfig(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
|
||||
@@ -1348,6 +1376,16 @@ func TestValidateConfigErrors(t *testing.T) {
|
||||
mutate: func(c *Config) { c.Gateway.OpenAIResponseHeaderTimeout = -1 },
|
||||
wantErr: "gateway.openai_response_header_timeout",
|
||||
},
|
||||
{
|
||||
name: "gateway openai first output timeout below minimum",
|
||||
mutate: func(c *Config) { c.Gateway.OpenAIFirstOutputTimeoutSeconds = 29 },
|
||||
wantErr: "gateway.openai_first_output_timeout_seconds",
|
||||
},
|
||||
{
|
||||
name: "gateway openai high effort first output timeout too large",
|
||||
mutate: func(c *Config) { c.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds = 1801 },
|
||||
wantErr: "gateway.openai_high_effort_first_output_timeout_seconds",
|
||||
},
|
||||
{
|
||||
name: "gateway max idle conns",
|
||||
mutate: func(c *Config) { c.Gateway.MaxIdleConns = 0 },
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestOpenAIForwardMayFailoverOnlyAfterNonSemanticWrite(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
before := service.OpenAICompactKeepaliveAdjustedWrittenSize(c)
|
||||
|
||||
_, err := fmt.Fprint(c.Writer, ":\n\n")
|
||||
require.NoError(t, err)
|
||||
c.Writer.Flush()
|
||||
|
||||
require.True(t, openAIForwardMayFailover(c, before, &service.UpstreamFailoverError{
|
||||
SafeToFailoverAfterWrite: true,
|
||||
}))
|
||||
require.False(t, openAIForwardMayFailover(c, before, &service.UpstreamFailoverError{}))
|
||||
}
|
||||
|
||||
func TestOpenAIFirstOutputFailoverStopsAfterOneAccountSwitch(t *testing.T) {
|
||||
failoverErr := &service.UpstreamFailoverError{SafeToFailoverAfterWrite: true}
|
||||
count := 0
|
||||
|
||||
require.False(t, openAIFirstOutputFailoverExhausted(failoverErr, &count))
|
||||
require.Equal(t, 1, count)
|
||||
require.True(t, openAIFirstOutputFailoverExhausted(failoverErr, &count))
|
||||
require.Equal(t, 1, count)
|
||||
}
|
||||
|
||||
func TestOpenAIRequestAllowsFailoverReplayStopsCanceledClient(t *testing.T) {
|
||||
require.False(t, openAIRequestAllowsFailoverReplay(nil))
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
requestCtx, cancel := context.WithCancel(context.Background())
|
||||
c.Request = httptest.NewRequest("POST", "/v1/responses", nil).WithContext(requestCtx)
|
||||
|
||||
require.True(t, openAIRequestAllowsFailoverReplay(c))
|
||||
cancel()
|
||||
require.False(t, openAIRequestAllowsFailoverReplay(c))
|
||||
}
|
||||
@@ -40,6 +40,8 @@ type OpenAIGatewayHandler struct {
|
||||
cfg *config.Config
|
||||
}
|
||||
|
||||
const maxOpenAIFirstOutputTimeoutSwitches = 1
|
||||
|
||||
func resolveOpenAIMessagesDispatchMappedModel(apiKey *service.APIKey, requestedModel string) string {
|
||||
if apiKey == nil || apiKey.Group == nil {
|
||||
return ""
|
||||
@@ -338,13 +340,17 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
|
||||
maxAccountSwitches := h.maxAccountSwitches
|
||||
switchCount := 0
|
||||
firstOutputTimeoutSwitchCount := 0
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
sameAccountRetryCount := make(map[int64]int)
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
var oauth429FailoverState service.OpenAIOAuth429FailoverState
|
||||
|
||||
for {
|
||||
if failoverClientGone(c) {
|
||||
// Streaming Forward intentionally detaches the upstream request so usage can
|
||||
// be drained after a disconnect. Re-check the client context before every
|
||||
// account attempt so a canceled request never starts a failover replay.
|
||||
if !openAIRequestAllowsFailoverReplay(c) {
|
||||
return
|
||||
}
|
||||
// Select account supporting the requested model
|
||||
@@ -467,10 +473,13 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
)
|
||||
return
|
||||
}
|
||||
if service.OpenAICompactKeepaliveAdjustedWrittenSize(c) != writerSizeBeforeForward {
|
||||
if !openAIForwardMayFailover(c, writerSizeBeforeForward, failoverErr) {
|
||||
h.handleFailoverExhausted(c, failoverErr, true)
|
||||
return
|
||||
}
|
||||
if failoverErr.SafeToFailoverAfterWrite && c.Writer.Written() {
|
||||
streamStarted = true
|
||||
}
|
||||
if failoverErr.ShouldReportAccountScheduleFailure() {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
}
|
||||
@@ -478,6 +487,10 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
h.handleFailoverExhausted(c, failoverErr, streamStarted)
|
||||
return
|
||||
}
|
||||
if openAIFirstOutputFailoverExhausted(failoverErr, &firstOutputTimeoutSwitchCount) {
|
||||
h.handleFailoverExhausted(c, failoverErr, streamStarted)
|
||||
return
|
||||
}
|
||||
// 池模式:同账号重试
|
||||
if failoverErr.RetryableOnSameAccount {
|
||||
retryLimit := account.GetPoolModeRetryCount()
|
||||
@@ -2284,6 +2297,34 @@ func openAIForwardErrorAlreadyCommunicated(c *gin.Context, writerSizeBeforeForwa
|
||||
return false
|
||||
}
|
||||
|
||||
func openAIForwardMayFailover(c *gin.Context, writerSizeBeforeForward int, failoverErr *service.UpstreamFailoverError) bool {
|
||||
if c == nil || c.Writer == nil {
|
||||
return false
|
||||
}
|
||||
if service.OpenAICompactKeepaliveAdjustedWrittenSize(c) == writerSizeBeforeForward {
|
||||
return true
|
||||
}
|
||||
return failoverErr != nil && failoverErr.SafeToFailoverAfterWrite
|
||||
}
|
||||
|
||||
func openAIRequestAllowsFailoverReplay(c *gin.Context) bool {
|
||||
if c == nil || c.Request == nil {
|
||||
return false
|
||||
}
|
||||
return !failoverClientGone(c)
|
||||
}
|
||||
|
||||
func openAIFirstOutputFailoverExhausted(failoverErr *service.UpstreamFailoverError, switchCount *int) bool {
|
||||
if failoverErr == nil || !failoverErr.SafeToFailoverAfterWrite || switchCount == nil {
|
||||
return false
|
||||
}
|
||||
if *switchCount >= maxOpenAIFirstOutputTimeoutSwitches {
|
||||
return true
|
||||
}
|
||||
*switchCount = *switchCount + 1
|
||||
return false
|
||||
}
|
||||
|
||||
// errorResponse returns OpenAI API format error response
|
||||
func (h *OpenAIGatewayHandler) errorResponse(c *gin.Context, status int, errType, message string) {
|
||||
// body-signal compact 心跳可能已把响应头提交为 200:JSON 错误体会与已
|
||||
|
||||
@@ -593,17 +593,18 @@ type GatewayFailureReason string
|
||||
// trigger account failover. Additive metadata keeps existing composite literals
|
||||
// source-compatible and preserves their legacy retry-next-account behavior.
|
||||
type UpstreamFailoverError struct {
|
||||
StatusCode int
|
||||
ResponseBody []byte // 上游响应体,用于错误透传规则匹配
|
||||
ResponseHeaders http.Header // 上游响应头,用于透传 cf-ray/cf-mitigated/content-type 等诊断信息
|
||||
ForceCacheBilling bool // Antigravity 粘性会话切换时设为 true
|
||||
RetryableOnSameAccount bool // 临时性错误(如 Google 间歇性 400、空响应),应在同一账号上重试 N 次再切换
|
||||
Stage GatewayFailureStage
|
||||
Scope GatewayFailureScope
|
||||
Reason GatewayFailureReason
|
||||
NextAccountAction NextAccountAction
|
||||
ClientStatusCode int
|
||||
ClientMessage string
|
||||
StatusCode int
|
||||
ResponseBody []byte // 上游响应体,用于错误透传规则匹配
|
||||
ResponseHeaders http.Header // 上游响应头,用于透传 cf-ray/cf-mitigated/content-type 等诊断信息
|
||||
ForceCacheBilling bool // Antigravity 粘性会话切换时设为 true
|
||||
RetryableOnSameAccount bool // 临时性错误(如 Google 间歇性 400、空响应),应在同一账号上重试 N 次再切换
|
||||
SafeToFailoverAfterWrite bool // 仅写出 SSE 注释等非语义字节时,仍可在同一客户端流中切换账号
|
||||
Stage GatewayFailureStage
|
||||
Scope GatewayFailureScope
|
||||
Reason GatewayFailureReason
|
||||
NextAccountAction NextAccountAction
|
||||
ClientStatusCode int
|
||||
ClientMessage string
|
||||
}
|
||||
|
||||
func (e *UpstreamFailoverError) Error() string {
|
||||
|
||||
@@ -0,0 +1,335 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const (
|
||||
openAIFirstOutputStageMemoryLimit = 64 * 1024
|
||||
openAIFirstOutputStageMaxBytes = 8 * 1024 * 1024
|
||||
openAIFirstOutputScannerFramingAllowance = 64
|
||||
openAIFirstOutputGuardQueueSize = 1
|
||||
openAIDefaultStreamQueueSize = 16
|
||||
)
|
||||
|
||||
var (
|
||||
errOpenAIFirstOutputStageLimit = errors.New("openai first-output staging limit exceeded")
|
||||
errOpenAIFirstOutputScannerLimit = errors.New("openai pre-output scanner token limit exceeded")
|
||||
)
|
||||
|
||||
type openAIFirstOutputStage struct {
|
||||
limit int64
|
||||
size int64
|
||||
memory bytes.Buffer
|
||||
tempFile *os.File
|
||||
tempPath string
|
||||
createTemp func() (*os.File, error)
|
||||
removeFile func(string) error
|
||||
memoryOnly bool
|
||||
cleanupErr error
|
||||
closed bool
|
||||
}
|
||||
|
||||
func newOpenAIFirstOutputStage(limit int64) *openAIFirstOutputStage {
|
||||
if limit < 1 {
|
||||
limit = 1
|
||||
}
|
||||
return &openAIFirstOutputStage{
|
||||
limit: limit,
|
||||
createTemp: func() (*os.File, error) { return os.CreateTemp("", "sub2api-openai-first-output-*") },
|
||||
removeFile: os.Remove,
|
||||
memoryOnly: runtime.GOOS == "windows",
|
||||
}
|
||||
}
|
||||
|
||||
func newDefaultOpenAIFirstOutputStage() *openAIFirstOutputStage {
|
||||
return newOpenAIFirstOutputStage(openAIFirstOutputStageMaxBytes)
|
||||
}
|
||||
|
||||
func openAIFirstOutputEventQueueSize(guardFirstOutput bool) int {
|
||||
if guardFirstOutput {
|
||||
return openAIFirstOutputGuardQueueSize
|
||||
}
|
||||
return openAIDefaultStreamQueueSize
|
||||
}
|
||||
|
||||
func openAIFirstOutputDynamicScanLines(guardActive *atomic.Bool) bufio.SplitFunc {
|
||||
return func(data []byte, atEOF bool) (advance int, token []byte, err error) {
|
||||
advance, token, err = bufio.ScanLines(data, atEOF)
|
||||
if err != nil || guardActive == nil || !guardActive.Load() {
|
||||
return advance, token, err
|
||||
}
|
||||
limit := openAIFirstOutputStageMaxBytes + openAIFirstOutputScannerFramingAllowance
|
||||
if token != nil {
|
||||
if len(token) > limit {
|
||||
return 0, nil, errOpenAIFirstOutputScannerLimit
|
||||
}
|
||||
return advance, token, nil
|
||||
}
|
||||
// At the limit with no delimiter, another byte would necessarily exceed
|
||||
// the guarded token budget. Fail before Scanner grows toward MaxLineSize.
|
||||
if len(data) >= limit {
|
||||
return 0, nil, errOpenAIFirstOutputScannerLimit
|
||||
}
|
||||
return advance, token, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (s *openAIFirstOutputStage) Buffered() int64 {
|
||||
if s == nil {
|
||||
return 0
|
||||
}
|
||||
return s.size
|
||||
}
|
||||
|
||||
func (s *openAIFirstOutputStage) WriteString(value string) (int, error) {
|
||||
if err := s.prepareWrite(len(value)); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
var n int
|
||||
var err error
|
||||
if s.tempFile == nil {
|
||||
n, err = s.memory.WriteString(value)
|
||||
} else {
|
||||
n, err = io.WriteString(s.tempFile, value)
|
||||
}
|
||||
s.size += int64(n)
|
||||
if err != nil {
|
||||
return n, fmt.Errorf("write first-output stage: %w", err)
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (s *openAIFirstOutputStage) Write(p []byte) (int, error) {
|
||||
if err := s.prepareWrite(len(p)); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
var n int
|
||||
var err error
|
||||
if s.tempFile == nil {
|
||||
n, err = s.memory.Write(p)
|
||||
} else {
|
||||
n, err = s.tempFile.Write(p)
|
||||
}
|
||||
s.size += int64(n)
|
||||
if err != nil {
|
||||
return n, fmt.Errorf("write first-output stage: %w", err)
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (s *openAIFirstOutputStage) prepareWrite(incoming int) error {
|
||||
if s == nil || s.closed {
|
||||
return os.ErrClosed
|
||||
}
|
||||
if int64(incoming) > s.limit-s.size {
|
||||
return fmt.Errorf("%w: buffered=%d incoming=%d limit=%d", errOpenAIFirstOutputStageLimit, s.size, incoming, s.limit)
|
||||
}
|
||||
if s.tempFile != nil || s.memoryOnly || s.size+int64(incoming) <= openAIFirstOutputStageMemoryLimit {
|
||||
return nil
|
||||
}
|
||||
file, err := s.createTemp()
|
||||
if err != nil {
|
||||
return fmt.Errorf("create first-output spool: %w", err)
|
||||
}
|
||||
path := file.Name()
|
||||
// Unlink before writing any request data. Unix keeps the file descriptor
|
||||
// readable, while crashes and SIGKILL cannot leave a named plaintext spool.
|
||||
if unlinkErr := s.removeFile(path); unlinkErr != nil {
|
||||
closeErr := file.Close()
|
||||
removeErr := s.removeFile(path)
|
||||
if errors.Is(removeErr, os.ErrNotExist) {
|
||||
removeErr = nil
|
||||
}
|
||||
s.memoryOnly = true
|
||||
if removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) {
|
||||
s.tempPath = path
|
||||
}
|
||||
s.cleanupErr = errors.Join(
|
||||
s.cleanupErr,
|
||||
fmt.Errorf("unlink first-output spool before use: %w", unlinkErr),
|
||||
closeErr,
|
||||
removeErr,
|
||||
)
|
||||
return nil
|
||||
}
|
||||
if _, err := file.Write(s.memory.Bytes()); err != nil {
|
||||
_ = file.Close()
|
||||
return fmt.Errorf("initialize first-output spool: %w", err)
|
||||
}
|
||||
s.tempFile = file
|
||||
s.tempPath = path
|
||||
s.memory.Reset()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *openAIFirstOutputStage) CommitTo(dst io.Writer) error {
|
||||
if s == nil || s.closed {
|
||||
return os.ErrClosed
|
||||
}
|
||||
if s.tempFile == nil {
|
||||
if _, err := io.Copy(dst, bytes.NewReader(s.memory.Bytes())); err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
if _, err := s.tempFile.Seek(0, io.SeekStart); err != nil {
|
||||
return fmt.Errorf("seek first-output spool: %w", err)
|
||||
}
|
||||
if _, err := io.CopyN(dst, s.tempFile, s.size); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := s.Close(); err != nil {
|
||||
// Delivery succeeded. Preserve cleanup failures for the handler's deferred
|
||||
// cleanup/logging pass instead of turning committed bytes into a stream error.
|
||||
s.cleanupErr = errors.Join(s.cleanupErr, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *openAIFirstOutputStage) Close() error {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
if s.closed && s.tempFile == nil && s.tempPath == "" && s.cleanupErr == nil {
|
||||
return nil
|
||||
}
|
||||
s.closed = true
|
||||
s.size = 0
|
||||
s.memory.Reset()
|
||||
closeErr := s.cleanupErr
|
||||
s.cleanupErr = nil
|
||||
if s.tempFile != nil {
|
||||
closeErr = errors.Join(closeErr, s.tempFile.Close())
|
||||
s.tempFile = nil
|
||||
}
|
||||
if s.tempPath != "" {
|
||||
removeErr := s.removeFile(s.tempPath)
|
||||
if removeErr == nil || errors.Is(removeErr, os.ErrNotExist) {
|
||||
s.tempPath = ""
|
||||
} else {
|
||||
closeErr = errors.Join(closeErr, removeErr)
|
||||
}
|
||||
}
|
||||
return closeErr
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) openAIFirstOutputTimeout(reasoningEffort string) time.Duration {
|
||||
if s == nil || s.cfg == nil || s.cfg.Gateway.OpenAIFirstOutputTimeoutSeconds <= 0 {
|
||||
return 0
|
||||
}
|
||||
seconds := s.cfg.Gateway.OpenAIFirstOutputTimeoutSeconds
|
||||
switch strings.ToLower(strings.TrimSpace(reasoningEffort)) {
|
||||
case "high", "xhigh", "max":
|
||||
if override := s.cfg.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds; override > 0 {
|
||||
seconds = override
|
||||
}
|
||||
}
|
||||
return time.Duration(seconds) * time.Second
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) newOpenAIFirstOutputTimeoutError(
|
||||
ctx context.Context,
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
startTime time.Time,
|
||||
originalModel string,
|
||||
reasoningEffort string,
|
||||
timeout time.Duration,
|
||||
phase string,
|
||||
responseHeaders http.Header,
|
||||
) *UpstreamFailoverError {
|
||||
elapsed := time.Since(startTime)
|
||||
logger.LegacyPrintf(
|
||||
"service.openai_gateway",
|
||||
"OpenAI first output timeout: account=%d model=%s effort=%s phase=%s elapsed=%s limit=%s",
|
||||
account.ID, originalModel, reasoningEffort, phase, elapsed, timeout,
|
||||
)
|
||||
requestID := strings.TrimSpace(responseHeaders.Get("x-request-id"))
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform, AccountID: account.ID, AccountName: account.Name,
|
||||
UpstreamStatusCode: http.StatusGatewayTimeout, UpstreamRequestID: requestID,
|
||||
Kind: "first_output_timeout", Message: "OpenAI upstream produced no semantic output before the deadline",
|
||||
Detail: fmt.Sprintf("phase=%s elapsed_ms=%d timeout_ms=%d", phase, elapsed.Milliseconds(), timeout.Milliseconds()),
|
||||
})
|
||||
if s.rateLimitService != nil {
|
||||
s.rateLimitService.HandleStreamTimeout(ctx, account, originalModel)
|
||||
}
|
||||
return &UpstreamFailoverError{
|
||||
StatusCode: http.StatusGatewayTimeout,
|
||||
ResponseBody: []byte(`{"error":{"type":"first_output_timeout","message":"Upstream produced no output before the deadline"}}`),
|
||||
ResponseHeaders: responseHeaders.Clone(), SafeToFailoverAfterWrite: true,
|
||||
}
|
||||
}
|
||||
|
||||
type openAIFirstOutputHeaderGuard struct {
|
||||
cancel context.CancelFunc
|
||||
release context.CancelFunc
|
||||
timer *time.Timer
|
||||
fired chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func newOpenAIFirstOutputHeaderGuard(
|
||||
ctx context.Context,
|
||||
release context.CancelFunc,
|
||||
deadline time.Time,
|
||||
) (context.Context, *openAIFirstOutputHeaderGuard) {
|
||||
guardedCtx, cancel := context.WithCancel(ctx)
|
||||
guard := &openAIFirstOutputHeaderGuard{cancel: cancel, release: release, fired: make(chan struct{})}
|
||||
remaining := time.Until(deadline)
|
||||
if remaining <= 0 {
|
||||
remaining = time.Nanosecond
|
||||
}
|
||||
guard.timer = time.AfterFunc(remaining, func() {
|
||||
close(guard.fired)
|
||||
cancel()
|
||||
})
|
||||
return guardedCtx, guard
|
||||
}
|
||||
|
||||
func (g *openAIFirstOutputHeaderGuard) stopHeaderWait() bool {
|
||||
if g.timer.Stop() {
|
||||
return false
|
||||
}
|
||||
<-g.fired
|
||||
return true
|
||||
}
|
||||
|
||||
func (g *openAIFirstOutputHeaderGuard) close() {
|
||||
g.once.Do(func() {
|
||||
g.timer.Stop()
|
||||
g.cancel()
|
||||
g.release()
|
||||
})
|
||||
}
|
||||
|
||||
type openAIRequestContextReadCloser struct {
|
||||
io.ReadCloser
|
||||
cleanup func()
|
||||
once sync.Once
|
||||
err error
|
||||
}
|
||||
|
||||
func (r *openAIRequestContextReadCloser) Close() error {
|
||||
r.once.Do(func() {
|
||||
r.cleanup()
|
||||
r.err = r.ReadCloser.Close()
|
||||
})
|
||||
return r.err
|
||||
}
|
||||
@@ -0,0 +1,640 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type blockingOpenAIResponseHeaderUpstream struct {
|
||||
canceled chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
type firstOutputCloseTrackingBody struct {
|
||||
io.ReadCloser
|
||||
closed chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (b *firstOutputCloseTrackingBody) Close() error {
|
||||
b.once.Do(func() { close(b.closed) })
|
||||
return b.ReadCloser.Close()
|
||||
}
|
||||
|
||||
func (u *blockingOpenAIResponseHeaderUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
select {
|
||||
case <-req.Context().Done():
|
||||
u.once.Do(func() { close(u.canceled) })
|
||||
return nil, req.Context().Err()
|
||||
case <-time.After(1500 * time.Millisecond):
|
||||
return nil, errors.New("test upstream was not canceled before response headers")
|
||||
}
|
||||
}
|
||||
|
||||
func (u *blockingOpenAIResponseHeaderUpstream) DoWithTLS(req *http.Request, _ string, _ int64, _ int, _ *tlsfingerprint.Profile) (*http.Response, error) {
|
||||
return u.Do(req, "", 0, 0)
|
||||
}
|
||||
|
||||
func TestOpenAIForwardFirstOutputTimeoutIncludesResponseHeaderWait(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
upstream := &blockingOpenAIResponseHeaderUpstream{canceled: make(chan struct{})}
|
||||
svc := &OpenAIGatewayService{
|
||||
cfg: &config.Config{Gateway: config.GatewayConfig{
|
||||
OpenAIFirstOutputTimeoutSeconds: 1,
|
||||
MaxLineSize: defaultMaxLineSize,
|
||||
}},
|
||||
httpUpstream: upstream,
|
||||
}
|
||||
body := []byte(`{"model":"gpt-5.5","stream":true,"reasoning":{"effort":"low"},"input":"hello"}`)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
|
||||
account := &Account{
|
||||
ID: 1, Name: "oauth-test", Platform: PlatformOpenAI, Type: AccountTypeOAuth,
|
||||
Status: StatusActive, Schedulable: true, Concurrency: 1,
|
||||
Credentials: map[string]any{"access_token": "test-token", "chatgpt_account_id": "test-account"},
|
||||
}
|
||||
|
||||
started := time.Now()
|
||||
_, err := svc.Forward(context.Background(), c, account, body)
|
||||
|
||||
require.Error(t, err)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, http.StatusGatewayTimeout, failoverErr.StatusCode)
|
||||
require.Contains(t, string(failoverErr.ResponseBody), "first_output_timeout")
|
||||
require.True(t, failoverErr.SafeToFailoverAfterWrite)
|
||||
require.Less(t, time.Since(started), 1300*time.Millisecond)
|
||||
require.Empty(t, rec.Body.String())
|
||||
select {
|
||||
case <-upstream.canceled:
|
||||
default:
|
||||
t.Fatal("response-header timeout did not cancel the upstream request context")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAINativeFirstOutputTimeoutDisabledPreservesSynchronousStream(t *testing.T) {
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{
|
||||
OpenAIFirstOutputTimeoutSeconds: 0,
|
||||
MaxLineSize: defaultMaxLineSize,
|
||||
}}}
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||||
`data: {"type":"response.created","response":{"id":"resp_disabled"}}`,
|
||||
"",
|
||||
`data: {"type":"response.completed","response":{"id":"resp_disabled","usage":{"input_tokens":1,"output_tokens":1}}}`,
|
||||
"",
|
||||
}, "\n")))}
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI}, time.Now(), "model", "model")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Contains(t, rec.Body.String(), "response.completed")
|
||||
}
|
||||
|
||||
func TestOpenAINativeFirstOutputTimeoutIgnoresPreambleAndCleansReader(t *testing.T) {
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{
|
||||
OpenAIFirstOutputTimeoutSeconds: 1,
|
||||
MaxLineSize: defaultMaxLineSize,
|
||||
}}}
|
||||
pr, pw := io.Pipe()
|
||||
writerDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(writerDone)
|
||||
defer func() { _ = pw.Close() }()
|
||||
_, _ = pw.Write([]byte("data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_slow\"}}\n\n"))
|
||||
_, _ = pw.Write([]byte("data: {\"type\":\"response.in_progress\",\"response\":{\"id\":\"resp_slow\"}}\n\n"))
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
}()
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
body := &firstOutputCloseTrackingBody{ReadCloser: pr, closed: make(chan struct{})}
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: body}
|
||||
|
||||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI}, time.Now().Add(-2*time.Second), "model", "model")
|
||||
|
||||
require.Error(t, err)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, http.StatusGatewayTimeout, failoverErr.StatusCode)
|
||||
require.Contains(t, string(failoverErr.ResponseBody), "first_output_timeout")
|
||||
require.True(t, failoverErr.SafeToFailoverAfterWrite)
|
||||
require.Empty(t, rec.Body.String())
|
||||
select {
|
||||
case <-body.closed:
|
||||
default:
|
||||
t.Fatal("first-output timeout did not close the upstream response body")
|
||||
}
|
||||
select {
|
||||
case <-writerDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("stream reader/writer goroutine did not exit after first-output timeout")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIFirstOutputTimeoutForReasoningEffort(t *testing.T) {
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{
|
||||
OpenAIFirstOutputTimeoutSeconds: 120,
|
||||
OpenAIHighEffortFirstOutputTimeoutSeconds: 300,
|
||||
}}}
|
||||
|
||||
require.Equal(t, 120*time.Second, svc.openAIFirstOutputTimeout("low"))
|
||||
require.Equal(t, 300*time.Second, svc.openAIFirstOutputTimeout("high"))
|
||||
require.Equal(t, 300*time.Second, svc.openAIFirstOutputTimeout("xhigh"))
|
||||
require.Equal(t, 300*time.Second, svc.openAIFirstOutputTimeout("max"))
|
||||
}
|
||||
|
||||
func TestOpenAIFirstOutputStageDefaultLimitIsIndependentFromScannerLimit(t *testing.T) {
|
||||
stage := newDefaultOpenAIFirstOutputStage()
|
||||
defer func() { require.NoError(t, stage.Close()) }()
|
||||
|
||||
require.EqualValues(t, 8*1024*1024, stage.limit)
|
||||
require.Greater(t, stage.limit, int64(68106))
|
||||
require.Less(t, stage.limit, int64(defaultMaxLineSize))
|
||||
}
|
||||
|
||||
func TestOpenAIFirstOutputEventQueueSizeBackpressuresGuardedStreams(t *testing.T) {
|
||||
require.Equal(t, 1, openAIFirstOutputEventQueueSize(true))
|
||||
require.Equal(t, 16, openAIFirstOutputEventQueueSize(false))
|
||||
}
|
||||
|
||||
func TestOpenAIFirstOutputDynamicScannerLimitsOnlyWhileGuardIsActive(t *testing.T) {
|
||||
var guardActive atomic.Bool
|
||||
guardActive.Store(true)
|
||||
split := openAIFirstOutputDynamicScanLines(&guardActive)
|
||||
guardLimit := openAIFirstOutputStageMaxBytes + openAIFirstOutputScannerFramingAllowance
|
||||
undelimited := bytes.Repeat([]byte("x"), guardLimit)
|
||||
|
||||
_, _, err := split(undelimited, false)
|
||||
require.ErrorIs(t, err, errOpenAIFirstOutputScannerLimit)
|
||||
|
||||
guardActive.Store(false)
|
||||
advance, token, err := split(undelimited, false)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, advance)
|
||||
require.Nil(t, token)
|
||||
}
|
||||
|
||||
func TestOpenAIFirstOutputStageOverflowIsAtomicAndCleanupRemovesSpool(t *testing.T) {
|
||||
stage := newOpenAIFirstOutputStage(70 * 1024)
|
||||
payload := bytes.Repeat([]byte("x"), 68*1024)
|
||||
n, err := stage.Write(payload)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, len(payload), n)
|
||||
if runtime.GOOS == "windows" {
|
||||
require.Nil(t, stage.tempFile)
|
||||
require.Empty(t, stage.tempPath)
|
||||
} else {
|
||||
require.NotNil(t, stage.tempFile)
|
||||
require.NotEmpty(t, stage.tempPath)
|
||||
_, err = os.Stat(stage.tempPath)
|
||||
require.ErrorIs(t, err, os.ErrNotExist)
|
||||
stat, statErr := stage.tempFile.Stat()
|
||||
require.NoError(t, statErr)
|
||||
require.Equal(t, os.FileMode(0o600), stat.Mode().Perm())
|
||||
}
|
||||
|
||||
n, err = stage.Write(bytes.Repeat([]byte("y"), 3*1024))
|
||||
require.Zero(t, n)
|
||||
require.ErrorIs(t, err, errOpenAIFirstOutputStageLimit)
|
||||
require.EqualValues(t, len(payload), stage.Buffered())
|
||||
path := stage.tempPath
|
||||
require.NoError(t, stage.Close())
|
||||
require.True(t, stage.closed)
|
||||
require.Nil(t, stage.tempFile)
|
||||
require.Empty(t, stage.tempPath)
|
||||
if path != "" {
|
||||
_, err = os.Stat(path)
|
||||
require.ErrorIs(t, err, os.ErrNotExist)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIFirstOutputStageCommitCopiesSpoolAndRemovesTemp(t *testing.T) {
|
||||
stage := newOpenAIFirstOutputStage(80 * 1024)
|
||||
payload := bytes.Repeat([]byte("z"), 68*1024)
|
||||
_, err := stage.Write(payload)
|
||||
require.NoError(t, err)
|
||||
path := stage.tempPath
|
||||
if runtime.GOOS == "windows" {
|
||||
require.Empty(t, path)
|
||||
require.Nil(t, stage.tempFile)
|
||||
} else {
|
||||
require.NotEmpty(t, path)
|
||||
require.NotNil(t, stage.tempFile)
|
||||
_, statErr := os.Stat(path)
|
||||
require.ErrorIs(t, statErr, os.ErrNotExist)
|
||||
}
|
||||
|
||||
var downstream bytes.Buffer
|
||||
require.NoError(t, stage.CommitTo(&downstream))
|
||||
require.Equal(t, payload, downstream.Bytes())
|
||||
require.Zero(t, stage.Buffered())
|
||||
if path != "" {
|
||||
_, err = os.Stat(path)
|
||||
require.ErrorIs(t, err, os.ErrNotExist)
|
||||
}
|
||||
require.NoError(t, stage.Close())
|
||||
}
|
||||
|
||||
func TestOpenAIFirstOutputStageUnlinkFailurePermanentlyFallsBackToMemoryAndRetriesCleanup(t *testing.T) {
|
||||
stage := newDefaultOpenAIFirstOutputStage()
|
||||
stage.memoryOnly = false
|
||||
t.Cleanup(func() {
|
||||
stage.removeFile = os.Remove
|
||||
_ = stage.Close()
|
||||
})
|
||||
createCalls := 0
|
||||
stage.createTemp = func() (*os.File, error) {
|
||||
createCalls++
|
||||
return os.CreateTemp("", "sub2api-openai-first-output-fallback-*")
|
||||
}
|
||||
removeCalls := 0
|
||||
stage.removeFile = func(path string) error {
|
||||
removeCalls++
|
||||
if removeCalls <= 2 {
|
||||
return errors.New("forced remove failure")
|
||||
}
|
||||
return os.Remove(path)
|
||||
}
|
||||
|
||||
payload := bytes.Repeat([]byte("m"), 68*1024)
|
||||
_, err := stage.Write(payload)
|
||||
require.NoError(t, err)
|
||||
require.True(t, stage.memoryOnly)
|
||||
require.Nil(t, stage.tempFile)
|
||||
require.NotEmpty(t, stage.tempPath)
|
||||
require.Equal(t, 1, createCalls)
|
||||
stat, statErr := os.Stat(stage.tempPath)
|
||||
require.NoError(t, statErr)
|
||||
require.Zero(t, stat.Size(), "failed-unlink fallback must never write plaintext to the named file")
|
||||
|
||||
_, err = stage.WriteString("more")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, createCalls, "memory-only fallback must not retry CreateTemp")
|
||||
path := stage.tempPath
|
||||
cleanupErr := stage.Close()
|
||||
require.ErrorContains(t, cleanupErr, "forced remove failure")
|
||||
require.Empty(t, stage.tempPath)
|
||||
_, err = os.Stat(path)
|
||||
require.ErrorIs(t, err, os.ErrNotExist)
|
||||
require.NoError(t, stage.Close())
|
||||
}
|
||||
|
||||
func TestOpenAINativeFirstOutputTimeoutDisarmsAfterSemanticOutput(t *testing.T) {
|
||||
cfg := &config.Config{Gateway: config.GatewayConfig{
|
||||
OpenAIFirstOutputTimeoutSeconds: 1,
|
||||
MaxLineSize: defaultMaxLineSize,
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: cfg, responseHeaderFilter: compileResponseHeaderFilter(cfg)}
|
||||
pr, pw := io.Pipe()
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
_, _ = pw.Write([]byte("data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_ok\"}}\n\n"))
|
||||
_, _ = pw.Write([]byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"hello\"}\n\n"))
|
||||
time.Sleep(1100 * time.Millisecond)
|
||||
_, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_ok\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n"))
|
||||
}()
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{
|
||||
"X-Request-Id": []string{"request-winning"},
|
||||
"X-Ratelimit-Remaining-Requests": []string{"42"},
|
||||
}, Body: pr}
|
||||
|
||||
result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI}, time.Now(), "model", "model")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, result.firstTokenMs)
|
||||
require.Contains(t, rec.Body.String(), "response.output_text.delta")
|
||||
require.Contains(t, rec.Body.String(), "response.completed")
|
||||
require.Equal(t, "request-winning", rec.Result().Header.Get("X-Request-Id"))
|
||||
require.Equal(t, "42", rec.Result().Header.Get("X-Ratelimit-Remaining-Requests"))
|
||||
}
|
||||
|
||||
func TestOpenAINativeFirstOutputTimeoutWaitsForCompleteSemanticEvent(t *testing.T) {
|
||||
const lineSize = 68106
|
||||
prefix := `data: {"type":"response.output_text.delta","delta":"`
|
||||
suffix := `"}`
|
||||
line := prefix + strings.Repeat("x", lineSize-len(prefix)-len(suffix)) + suffix
|
||||
require.Len(t, line, lineSize)
|
||||
assertOpenAINativeLargeOpenEventTimesOutWithoutLeak(t, line)
|
||||
}
|
||||
|
||||
func TestOpenAINativeFirstOutputTimeoutDoesNotLeakLargePreambleEvent(t *testing.T) {
|
||||
const lineSize = 68106
|
||||
prefix := `data: {"type":"response.created","response":{"id":"resp_partial","padding":"`
|
||||
suffix := `"}}`
|
||||
line := prefix + strings.Repeat("x", lineSize-len(prefix)-len(suffix)) + suffix
|
||||
require.Len(t, line, lineSize)
|
||||
assertOpenAINativeLargeOpenEventTimesOutWithoutLeak(t, line)
|
||||
}
|
||||
|
||||
func assertOpenAINativeLargeOpenEventTimesOutWithoutLeak(t *testing.T, line string) {
|
||||
t.Helper()
|
||||
cfg := &config.Config{Gateway: config.GatewayConfig{
|
||||
OpenAIFirstOutputTimeoutSeconds: 1,
|
||||
StreamKeepaliveInterval: 1,
|
||||
MaxLineSize: defaultMaxLineSize,
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: cfg, responseHeaderFilter: compileResponseHeaderFilter(cfg)}
|
||||
pr, pw := io.Pipe()
|
||||
body := &firstOutputCloseTrackingBody{ReadCloser: pr, closed: make(chan struct{})}
|
||||
writerDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(writerDone)
|
||||
defer func() { _ = pw.Close() }()
|
||||
_, _ = pw.Write([]byte(line + "\n"))
|
||||
select {
|
||||
case <-body.closed:
|
||||
case <-time.After(2 * time.Second):
|
||||
}
|
||||
}()
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{
|
||||
"X-Request-Id": []string{"request-partial"},
|
||||
"X-Ratelimit-Remaining-Requests": []string{"1"},
|
||||
}, Body: body}
|
||||
|
||||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI}, time.Now(), "model", "model")
|
||||
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, http.StatusGatewayTimeout, failoverErr.StatusCode)
|
||||
require.Contains(t, string(failoverErr.ResponseBody), "first_output_timeout")
|
||||
require.True(t, failoverErr.SafeToFailoverAfterWrite)
|
||||
require.NotContains(t, rec.Body.String(), "data:", "attempt JSON must remain private before the SSE boundary")
|
||||
require.NotContains(t, rec.Body.String(), `"type"`, "attempt JSON must remain private before the SSE boundary")
|
||||
for _, outputLine := range strings.Split(strings.TrimSpace(rec.Body.String()), "\n") {
|
||||
if outputLine != "" {
|
||||
require.True(t, strings.HasPrefix(outputLine, ":"), "only keepalive comments may precede failover: %q", outputLine)
|
||||
}
|
||||
}
|
||||
require.Empty(t, rec.Header().Values("X-Request-Id"))
|
||||
require.Empty(t, rec.Header().Values("X-Ratelimit-Remaining-Requests"))
|
||||
select {
|
||||
case <-writerDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("partial-event writer did not exit after timeout closed the body")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAINativeFirstOutputEOFDispatchesTerminalEventWithoutBlankLine(t *testing.T) {
|
||||
cfg := &config.Config{Gateway: config.GatewayConfig{
|
||||
OpenAIFirstOutputTimeoutSeconds: 1,
|
||||
MaxLineSize: defaultMaxLineSize,
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: cfg, responseHeaderFilter: compileResponseHeaderFilter(cfg)}
|
||||
payload := `data: {"type":"response.completed","response":{"id":"resp_eof","usage":{"input_tokens":3,"output_tokens":2}}}`
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"X-Request-Id": []string{"request-eof"},
|
||||
"X-Ratelimit-Remaining-Requests": []string{"17"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(payload)),
|
||||
}
|
||||
|
||||
result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI}, time.Now(), "model", "model")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, result.firstTokenMs)
|
||||
require.Equal(t, "resp_eof", result.responseID)
|
||||
require.Equal(t, 3, result.usage.InputTokens)
|
||||
require.Equal(t, 2, result.usage.OutputTokens)
|
||||
require.Contains(t, rec.Body.String(), `"type":"response.completed"`)
|
||||
require.Contains(t, rec.Body.String(), `"id":"resp_eof"`)
|
||||
require.True(t, strings.HasSuffix(rec.Body.String(), "\n"))
|
||||
require.False(t, strings.HasSuffix(rec.Body.String(), "\n\n"), "EOF dispatch must not synthesize a blank line")
|
||||
require.Equal(t, "request-eof", rec.Result().Header.Get("X-Request-Id"))
|
||||
require.Equal(t, "17", rec.Result().Header.Get("X-Ratelimit-Remaining-Requests"))
|
||||
}
|
||||
|
||||
func TestOpenAINativeFirstOutputStageOverflowFailsOverWithoutAttemptBytes(t *testing.T) {
|
||||
cfg := &config.Config{Gateway: config.GatewayConfig{
|
||||
OpenAIFirstOutputTimeoutSeconds: 30,
|
||||
MaxLineSize: 2 * 1024 * 1024,
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: cfg, responseHeaderFilter: compileResponseHeaderFilter(cfg)}
|
||||
const lineSize = 1024*1024 - 256
|
||||
prefix := `data: {"type":"response.output_text.delta","delta":"`
|
||||
suffix := `"}`
|
||||
line := prefix + strings.Repeat("x", lineSize-len(prefix)-len(suffix)) + suffix
|
||||
body := strings.Repeat(line+"\n", 9)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"X-Request-Id": []string{"request-overflow"},
|
||||
"X-Ratelimit-Remaining-Requests": []string{"1"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
}
|
||||
|
||||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI}, time.Now(), "model", "model")
|
||||
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
|
||||
require.True(t, failoverErr.SafeToFailoverAfterWrite)
|
||||
require.Contains(t, string(failoverErr.ResponseBody), "staging limit exceeded")
|
||||
require.Empty(t, rec.Body.String())
|
||||
require.Empty(t, rec.Header().Values("X-Request-Id"))
|
||||
require.Empty(t, rec.Header().Values("X-Ratelimit-Remaining-Requests"))
|
||||
}
|
||||
|
||||
func TestOpenAINativeFirstOutputScannerRejectsOversizedLineWithoutLeak(t *testing.T) {
|
||||
cfg := &config.Config{Gateway: config.GatewayConfig{
|
||||
OpenAIFirstOutputTimeoutSeconds: 30,
|
||||
MaxLineSize: defaultMaxLineSize,
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: cfg, responseHeaderFilter: compileResponseHeaderFilter(cfg)}
|
||||
oversizedLine := "data: " + strings.Repeat("x", openAIFirstOutputStageMaxBytes+openAIFirstOutputScannerFramingAllowance+1024)
|
||||
body := "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_private\"}}\n\n" + oversizedLine + "\n"
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"X-Request-Id": []string{"request-too-large"},
|
||||
"X-Ratelimit-Remaining-Requests": []string{"1"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
}
|
||||
|
||||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI}, time.Now(), "model", "model")
|
||||
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
|
||||
require.True(t, failoverErr.SafeToFailoverAfterWrite)
|
||||
require.Contains(t, string(failoverErr.ResponseBody), "line exceeds guarded first-output limit")
|
||||
require.Empty(t, rec.Body.String())
|
||||
require.Empty(t, rec.Header().Values("X-Request-Id"))
|
||||
require.Empty(t, rec.Header().Values("X-Ratelimit-Remaining-Requests"))
|
||||
}
|
||||
|
||||
func TestOpenAINativeFirstOutputScannerAllowsLargeEventAfterSemanticBoundary(t *testing.T) {
|
||||
cfg := &config.Config{Gateway: config.GatewayConfig{
|
||||
OpenAIFirstOutputTimeoutSeconds: 30,
|
||||
MaxLineSize: defaultMaxLineSize,
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: cfg, responseHeaderFilter: compileResponseHeaderFilter(cfg)}
|
||||
largeDelta := strings.Repeat("i", openAIFirstOutputStageMaxBytes+openAIFirstOutputScannerFramingAllowance+1024)
|
||||
body := strings.Join([]string{
|
||||
`data: {"type":"response.output_text.delta","delta":"ready"}`,
|
||||
"",
|
||||
`data: {"type":"response.output_text.delta","delta":"` + largeDelta + `"}`,
|
||||
"",
|
||||
`data: {"type":"response.completed","response":{"id":"resp_large_image","usage":{"input_tokens":4,"output_tokens":3}}}`,
|
||||
"",
|
||||
}, "\n")
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"X-Request-Id": []string{"request-large-image"}},
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
}
|
||||
|
||||
result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI}, time.Now(), "model", "model")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, result.firstTokenMs)
|
||||
require.Equal(t, "resp_large_image", result.responseID)
|
||||
require.Equal(t, 4, result.usage.InputTokens)
|
||||
require.Equal(t, 3, result.usage.OutputTokens)
|
||||
require.Contains(t, rec.Body.String(), `"delta":"ready"`)
|
||||
require.Contains(t, rec.Body.String(), `"id":"resp_large_image"`)
|
||||
require.Contains(t, rec.Body.String(), strings.Repeat("i", 1024))
|
||||
require.Equal(t, "request-large-image", rec.Result().Header.Get("X-Request-Id"))
|
||||
}
|
||||
|
||||
func TestOpenAINativeFirstOutputTimeoutDisabledPreservesKeepaliveFlush(t *testing.T) {
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{
|
||||
StreamKeepaliveInterval: 1,
|
||||
MaxLineSize: defaultMaxLineSize,
|
||||
}}}
|
||||
pr, pw := io.Pipe()
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
_, _ = pw.Write([]byte("data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_stalled\"}}\n\n"))
|
||||
_, _ = pw.Write([]byte("data: {\"type\":\"response.in_progress\",\"response\":{\"id\":\"resp_stalled\"}}\n\n"))
|
||||
time.Sleep(1100 * time.Millisecond)
|
||||
}()
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: pr}
|
||||
|
||||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI}, time.Now(), "model", "model")
|
||||
|
||||
require.Error(t, err)
|
||||
require.Contains(t, rec.Body.String(), ":\n\n")
|
||||
require.Contains(t, rec.Body.String(), "response.created")
|
||||
require.Contains(t, rec.Body.String(), "response.in_progress")
|
||||
}
|
||||
|
||||
func TestOpenAINativeFirstOutputFailoverKeepsAttemptHeadersPrivateAfterKeepaliveCommit(t *testing.T) {
|
||||
cfg := &config.Config{Gateway: config.GatewayConfig{
|
||||
OpenAIFirstOutputTimeoutSeconds: 2,
|
||||
StreamKeepaliveInterval: 1,
|
||||
MaxLineSize: defaultMaxLineSize,
|
||||
}}
|
||||
svc := &OpenAIGatewayService{
|
||||
cfg: cfg,
|
||||
responseHeaderFilter: compileResponseHeaderFilter(cfg),
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
firstBody, firstWriter := io.Pipe()
|
||||
trackedFirstBody := &firstOutputCloseTrackingBody{ReadCloser: firstBody, closed: make(chan struct{})}
|
||||
firstWriterDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(firstWriterDone)
|
||||
defer func() { _ = firstWriter.Close() }()
|
||||
_, _ = firstWriter.Write([]byte("data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_first\"}}\n\n"))
|
||||
select {
|
||||
case <-trackedFirstBody.closed:
|
||||
case <-time.After(4 * time.Second):
|
||||
}
|
||||
}()
|
||||
firstResp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"text/event-stream"},
|
||||
"X-Request-Id": []string{"request-first"},
|
||||
"X-Ratelimit-Remaining-Requests": []string{"1"},
|
||||
},
|
||||
Body: trackedFirstBody,
|
||||
}
|
||||
|
||||
_, firstErr := svc.handleStreamingResponse(c.Request.Context(), firstResp, c, &Account{ID: 1, Platform: PlatformOpenAI}, time.Now(), "model", "model")
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, firstErr, &failoverErr)
|
||||
require.Contains(t, rec.Body.String(), ":\n\n", "first attempt should have committed only a stable keepalive")
|
||||
require.NotContains(t, rec.Body.String(), "resp_first")
|
||||
|
||||
secondResp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"text/event-stream"},
|
||||
"X-Request-Id": []string{"request-second"},
|
||||
"X-Ratelimit-Remaining-Requests": []string{"99"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||||
`data: {"type":"response.output_text.delta","delta":"hello"}`,
|
||||
"",
|
||||
`data: {"type":"response.completed","response":{"id":"resp_second","usage":{"input_tokens":1,"output_tokens":1}}}`,
|
||||
"",
|
||||
}, "\n"))),
|
||||
}
|
||||
result, secondErr := svc.handleStreamingResponse(c.Request.Context(), secondResp, c, &Account{ID: 2, Platform: PlatformOpenAI}, time.Now(), "model", "model")
|
||||
|
||||
require.NoError(t, secondErr)
|
||||
require.NotNil(t, result)
|
||||
require.Contains(t, rec.Body.String(), "resp_second")
|
||||
wireHeaders := rec.Result().Header
|
||||
require.Empty(t, wireHeaders.Values("X-Request-Id"))
|
||||
require.Empty(t, wireHeaders.Values("X-Ratelimit-Remaining-Requests"))
|
||||
require.Empty(t, rec.Header().Values("X-Request-Id"))
|
||||
require.Empty(t, rec.Header().Values("X-Ratelimit-Remaining-Requests"))
|
||||
select {
|
||||
case <-firstWriterDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("first account writer did not exit after timeout")
|
||||
}
|
||||
}
|
||||
@@ -730,14 +730,37 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
return nil, wsErr
|
||||
}
|
||||
|
||||
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, upstreamModel, billingModel, originalModel)
|
||||
// 国产模型默认 effort 补充:此处 reqModel 已被 mapping 重写为 billingModel。
|
||||
reasoningEffort = ApplyThinkingEnabledFallback(reasoningEffort, body, reqModel)
|
||||
reasoningEffortValue := ""
|
||||
if reasoningEffort != nil {
|
||||
reasoningEffortValue = *reasoningEffort
|
||||
}
|
||||
firstOutputTimeout := time.Duration(0)
|
||||
if reqStream && account.Platform == PlatformOpenAI {
|
||||
firstOutputTimeout = s.openAIFirstOutputTimeout(reasoningEffortValue)
|
||||
}
|
||||
|
||||
httpInvalidEncryptedContentRetryTried := false
|
||||
agentTaskRecoveryTried := false
|
||||
for {
|
||||
// Build upstream request
|
||||
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
|
||||
var headerGuard *openAIFirstOutputHeaderGuard
|
||||
if firstOutputTimeout > 0 {
|
||||
upstreamCtx, headerGuard = newOpenAIFirstOutputHeaderGuard(
|
||||
upstreamCtx, releaseUpstreamCtx, startTime.Add(firstOutputTimeout),
|
||||
)
|
||||
}
|
||||
upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, body, token, reqStream, promptCacheKey, isCodexCLI)
|
||||
releaseUpstreamCtx()
|
||||
if headerGuard == nil {
|
||||
releaseUpstreamCtx()
|
||||
}
|
||||
if err != nil {
|
||||
if headerGuard != nil {
|
||||
headerGuard.close()
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -751,12 +774,31 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
upstreamStart := time.Now()
|
||||
resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
|
||||
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
|
||||
if headerGuard != nil && headerGuard.stopHeaderWait() {
|
||||
if resp != nil && resp.Body != nil {
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
headerGuard.close()
|
||||
return nil, s.newOpenAIFirstOutputTimeoutError(
|
||||
ctx, c, account, startTime, originalModel, reasoningEffortValue,
|
||||
firstOutputTimeout, "response_headers", nil,
|
||||
)
|
||||
}
|
||||
if err != nil {
|
||||
if resp != nil && resp.Body != nil {
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
if headerGuard != nil {
|
||||
headerGuard.close()
|
||||
}
|
||||
// Transport-level failure (proxy/DNS/TCP/TLS — no HTTP response). Convert to
|
||||
// a failover so the handler switches to a healthy account, and temporarily
|
||||
// unschedule the account on durable faults (e.g. rejected proxy credentials).
|
||||
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
|
||||
}
|
||||
if headerGuard != nil {
|
||||
resp.Body = &openAIRequestContextReadCloser{ReadCloser: resp.Body, cleanup: headerGuard.close}
|
||||
}
|
||||
|
||||
// Handle error response
|
||||
if resp.StatusCode >= 400 {
|
||||
@@ -824,10 +866,6 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, upstreamModel, billingModel, originalModel)
|
||||
// 国产模型默认 effort 补充:此处 reqModel 已被 mapping 重写为 billingModel(见
|
||||
// line 2510-2515 的 GetMappedModel + reqModel 赋值),可直接作为 mappedModel。
|
||||
reasoningEffort = ApplyThinkingEnabledFallback(reasoningEffort, body, reqModel)
|
||||
serviceTier := extractOpenAIServiceTierFromBody(body)
|
||||
// 上游接受后只保留计费需要的标量,避免响应处理期间继续保活完整 input/tools map。
|
||||
reqBody = nil
|
||||
@@ -839,7 +877,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
imageCount := 0
|
||||
var imageOutputSizes []string
|
||||
if reqStream {
|
||||
streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, upstreamModel)
|
||||
streamResult, err := s.handleStreamingResponseWithReasoning(ctx, resp, c, account, startTime, originalModel, upstreamModel, reasoningEffortValue)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -39,7 +39,23 @@ type openaiNonStreamingResult struct {
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel, mappedModel string) (*openaiStreamingResult, error) {
|
||||
if s.responseHeaderFilter != nil {
|
||||
return s.handleStreamingResponseWithReasoning(ctx, resp, c, account, startTime, originalModel, mappedModel, "")
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel, mappedModel, reasoningEffort string) (*openaiStreamingResult, error) {
|
||||
firstOutputTimeout := time.Duration(0)
|
||||
if account != nil && account.Platform == PlatformOpenAI {
|
||||
firstOutputTimeout = s.openAIFirstOutputTimeout(reasoningEffort)
|
||||
}
|
||||
guardFirstOutput := firstOutputTimeout > 0
|
||||
var attemptResponseHeaders http.Header
|
||||
if guardFirstOutput {
|
||||
if s.responseHeaderFilter != nil {
|
||||
attemptResponseHeaders = responseheaders.FilterHeaders(resp.Header, s.responseHeaderFilter)
|
||||
} else if requestID := strings.TrimSpace(resp.Header.Get("x-request-id")); requestID != "" {
|
||||
attemptResponseHeaders = http.Header{"X-Request-Id": []string{requestID}}
|
||||
}
|
||||
} else if s.responseHeaderFilter != nil {
|
||||
responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
|
||||
}
|
||||
|
||||
@@ -50,19 +66,68 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
c.Header("X-Accel-Buffering", "no")
|
||||
|
||||
// Pass through other headers
|
||||
if v := resp.Header.Get("x-request-id"); v != "" {
|
||||
if !guardFirstOutput && resp.Header.Get("x-request-id") != "" {
|
||||
v := resp.Header.Get("x-request-id")
|
||||
c.Header("x-request-id", v)
|
||||
}
|
||||
applyAttemptResponseHeaders := func() {
|
||||
if !guardFirstOutput || len(attemptResponseHeaders) == 0 || c.Writer.Written() {
|
||||
return
|
||||
}
|
||||
for key, values := range attemptResponseHeaders {
|
||||
for _, value := range values {
|
||||
c.Writer.Header().Add(key, value)
|
||||
}
|
||||
}
|
||||
// These headers describe this gateway's SSE stream and are stable across
|
||||
// account attempts. Keep them authoritative over upstream values.
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("X-Accel-Buffering", "no")
|
||||
}
|
||||
|
||||
w := c.Writer
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
return nil, errors.New("streaming not supported")
|
||||
}
|
||||
maxLineSize := defaultMaxLineSize
|
||||
if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
|
||||
maxLineSize = s.cfg.Gateway.MaxLineSize
|
||||
}
|
||||
var firstTokenMs *int
|
||||
bufferedWriter := bufio.NewWriterSize(w, 4*1024)
|
||||
var firstOutputStage *openAIFirstOutputStage
|
||||
if guardFirstOutput {
|
||||
firstOutputStage = newDefaultOpenAIFirstOutputStage()
|
||||
defer func() {
|
||||
if err := firstOutputStage.Close(); err != nil {
|
||||
logger.LegacyPrintf("service.openai_gateway", "OpenAI first-output staging cleanup failed: account=%d model=%s error=%v", account.ID, originalModel, err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
writePendingString := func(value string) (int, error) {
|
||||
if firstOutputStage != nil && firstTokenMs == nil && !firstOutputStage.closed {
|
||||
return firstOutputStage.WriteString(value)
|
||||
}
|
||||
return bufferedWriter.WriteString(value)
|
||||
}
|
||||
pendingBytes := func() int64 {
|
||||
if firstOutputStage != nil && firstTokenMs == nil && !firstOutputStage.closed {
|
||||
return firstOutputStage.Buffered()
|
||||
}
|
||||
return int64(bufferedWriter.Buffered())
|
||||
}
|
||||
flushBuffered := func() error {
|
||||
if err := bufferedWriter.Flush(); err != nil {
|
||||
return err
|
||||
if firstOutputStage != nil && firstTokenMs == nil && !firstOutputStage.closed {
|
||||
if err := firstOutputStage.CommitTo(w); err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
if err := bufferedWriter.Flush(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
flusher.Flush()
|
||||
return nil
|
||||
@@ -70,15 +135,15 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
|
||||
usage := &OpenAIUsage{}
|
||||
imageCounter := newOpenAIImageOutputCounter()
|
||||
var firstTokenMs *int
|
||||
responseID := ""
|
||||
var firstOutputScanGuard atomic.Bool
|
||||
firstOutputScanGuard.Store(guardFirstOutput)
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
maxLineSize := defaultMaxLineSize
|
||||
if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
|
||||
maxLineSize = s.cfg.Gateway.MaxLineSize
|
||||
}
|
||||
scanBuf := getSSEScannerBuf64K()
|
||||
scanner.Buffer(scanBuf[:0], maxLineSize)
|
||||
if guardFirstOutput {
|
||||
scanner.Split(openAIFirstOutputDynamicScanLines(&firstOutputScanGuard))
|
||||
}
|
||||
documentScanner := newOpenAISSEJSONDocumentScanner(scanner)
|
||||
|
||||
streamInterval := time.Duration(0)
|
||||
@@ -110,6 +175,31 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
if keepaliveTicker != nil {
|
||||
keepaliveCh = keepaliveTicker.C
|
||||
}
|
||||
|
||||
var firstOutputTimer *time.Timer
|
||||
var firstOutputCh <-chan time.Time
|
||||
if firstOutputTimeout > 0 {
|
||||
remaining := time.Until(startTime.Add(firstOutputTimeout))
|
||||
if remaining <= 0 {
|
||||
remaining = time.Nanosecond
|
||||
}
|
||||
firstOutputTimer = time.NewTimer(remaining)
|
||||
firstOutputCh = firstOutputTimer.C
|
||||
defer firstOutputTimer.Stop()
|
||||
}
|
||||
stopFirstOutputTimer := func() {
|
||||
if firstOutputTimer == nil {
|
||||
return
|
||||
}
|
||||
if !firstOutputTimer.Stop() {
|
||||
select {
|
||||
case <-firstOutputTimer.C:
|
||||
default:
|
||||
}
|
||||
}
|
||||
firstOutputTimer = nil
|
||||
firstOutputCh = nil
|
||||
}
|
||||
// Track downstream writes separately from upstream reads: pre-output failover
|
||||
// can buffer response.created / response.in_progress, so keepalive must be
|
||||
// based on downstream idle time.
|
||||
@@ -126,8 +216,52 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
clientOutputStarted := false
|
||||
upstreamRequestID := strings.TrimSpace(resp.Header.Get("x-request-id"))
|
||||
var streamEarlyErr error
|
||||
eventShouldFlush := false
|
||||
eventInProgress := false
|
||||
eventStartsClientOutput := false
|
||||
eventShouldFlush := false
|
||||
handlePendingWriteError := func(err error) {
|
||||
if firstOutputStage != nil && firstTokenMs == nil && !firstOutputStage.closed {
|
||||
message := "OpenAI first-output staging failed"
|
||||
if errors.Is(err, errOpenAIFirstOutputStageLimit) {
|
||||
message = "OpenAI first-output staging limit exceeded"
|
||||
}
|
||||
logger.LegacyPrintf("service.openai_gateway", "%s: account=%d model=%s error=%v", message, account.ID, originalModel, err)
|
||||
failoverErr := s.newOpenAIStreamFailoverError(c, account, false, upstreamRequestID, nil, message)
|
||||
failoverErr.SafeToFailoverAfterWrite = true
|
||||
streamEarlyErr = failoverErr
|
||||
_ = resp.Body.Close()
|
||||
return
|
||||
}
|
||||
clientDisconnected = true
|
||||
logger.LegacyPrintf("service.openai_gateway", "Client disconnected during streaming, continuing to drain upstream for billing")
|
||||
}
|
||||
completeGuardedEvent := func(queueDrained bool) {
|
||||
completedSemanticEvent := eventStartsClientOutput
|
||||
shouldFlush := eventShouldFlush || (queueDrained && clientOutputStarted)
|
||||
eventInProgress = false
|
||||
if !clientDisconnected {
|
||||
if completedSemanticEvent {
|
||||
applyAttemptResponseHeaders()
|
||||
}
|
||||
if shouldFlush {
|
||||
if err := flushBuffered(); err != nil {
|
||||
clientDisconnected = true
|
||||
logger.LegacyPrintf("service.openai_gateway", "Client disconnected during streaming flush, continuing to drain upstream for billing")
|
||||
} else {
|
||||
clientOutputStarted = true
|
||||
lastDownstreamWriteAt = time.Now()
|
||||
}
|
||||
}
|
||||
}
|
||||
if completedSemanticEvent && firstTokenMs == nil {
|
||||
firstOutputScanGuard.Store(false)
|
||||
ms := int(time.Since(startTime).Milliseconds())
|
||||
firstTokenMs = &ms
|
||||
stopFirstOutputTimer()
|
||||
}
|
||||
eventStartsClientOutput = false
|
||||
eventShouldFlush = false
|
||||
}
|
||||
sendErrorEvent := func(reason string) {
|
||||
if errorEventSent || clientDisconnected {
|
||||
return
|
||||
@@ -138,7 +272,7 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
clientDisconnected = true
|
||||
return
|
||||
}
|
||||
if _, err := bufferedWriter.WriteString("data: " + payload + "\n\n"); err != nil {
|
||||
if _, err := writePendingString("data: " + payload + "\n\n"); err != nil {
|
||||
clientDisconnected = true
|
||||
return
|
||||
}
|
||||
@@ -164,7 +298,7 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
}
|
||||
}
|
||||
flushPending := func(disconnectMessage string) {
|
||||
if clientDisconnected || bufferedWriter.Buffered() == 0 {
|
||||
if clientDisconnected || pendingBytes() == 0 {
|
||||
return
|
||||
}
|
||||
if err := flushBuffered(); err != nil {
|
||||
@@ -176,6 +310,10 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
lastDownstreamWriteAt = time.Now()
|
||||
}
|
||||
finalizeStream := func() (*openaiStreamingResult, error) {
|
||||
if guardFirstOutput && eventInProgress {
|
||||
// EOF dispatches the final SSE event even without a trailing blank line.
|
||||
completeGuardedEvent(true)
|
||||
}
|
||||
if !sawTerminalEvent && !openAIStreamClientOutputStarted(c, clientOutputStarted) && !eventShouldFlush {
|
||||
return resultWithUsage(), s.newOpenAIStreamFailoverError(
|
||||
c,
|
||||
@@ -199,6 +337,24 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
if scanErr == nil {
|
||||
return nil, nil, false
|
||||
}
|
||||
if errors.Is(scanErr, errOpenAIFirstOutputScannerLimit) && firstTokenMs == nil {
|
||||
logger.LegacyPrintf("service.openai_gateway", "SSE token exceeded guarded first-output limit: account=%d limit=%d error=%v", account.ID, openAIFirstOutputStageMaxBytes+openAIFirstOutputScannerFramingAllowance, scanErr)
|
||||
failoverErr := s.newOpenAIStreamFailoverError(
|
||||
c, account, false, upstreamRequestID, nil,
|
||||
"OpenAI SSE line exceeds guarded first-output limit",
|
||||
)
|
||||
failoverErr.SafeToFailoverAfterWrite = true
|
||||
return resultWithUsage(), failoverErr, true
|
||||
}
|
||||
if errors.Is(scanErr, bufio.ErrTooLong) && guardFirstOutput && firstTokenMs == nil {
|
||||
logger.LegacyPrintf("service.openai_gateway", "SSE line too long before first output: account=%d max_size=%d error=%v", account.ID, maxLineSize, scanErr)
|
||||
failoverErr := s.newOpenAIStreamFailoverError(
|
||||
c, account, false, upstreamRequestID, nil,
|
||||
"OpenAI SSE line exceeds guarded first-output limit",
|
||||
)
|
||||
failoverErr.SafeToFailoverAfterWrite = true
|
||||
return resultWithUsage(), failoverErr, true
|
||||
}
|
||||
if sawTerminalEvent {
|
||||
if !sawFailedEvent {
|
||||
logger.LegacyPrintf("service.openai_gateway", "Upstream scan ended after terminal event: %v", scanErr)
|
||||
@@ -345,6 +501,9 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
line = s.replaceModelInSSELine(line, mappedModel, originalModel)
|
||||
}
|
||||
startsClientOutput := forceFlushFailedEvent || openAIStreamDataStartsClientOutput(data, eventType)
|
||||
if guardFirstOutput {
|
||||
eventStartsClientOutput = eventStartsClientOutput || startsClientOutput
|
||||
}
|
||||
|
||||
// 写入客户端(客户端断开后继续 drain 上游)
|
||||
if !clientDisconnected {
|
||||
@@ -354,40 +513,49 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
shouldFlush = true
|
||||
}
|
||||
eventShouldFlush = eventShouldFlush || shouldFlush
|
||||
if _, err := bufferedWriter.WriteString(line); err != nil {
|
||||
clientDisconnected = true
|
||||
logger.LegacyPrintf("service.openai_gateway", "Client disconnected during streaming, continuing to drain upstream for billing")
|
||||
} else if _, err := bufferedWriter.WriteString("\n"); err != nil {
|
||||
clientDisconnected = true
|
||||
logger.LegacyPrintf("service.openai_gateway", "Client disconnected during streaming, continuing to drain upstream for billing")
|
||||
if _, err := writePendingString(line); err != nil {
|
||||
handlePendingWriteError(err)
|
||||
} else if _, err := writePendingString("\n"); err != nil {
|
||||
handlePendingWriteError(err)
|
||||
} else {
|
||||
eventInProgress = true
|
||||
}
|
||||
}
|
||||
|
||||
// Record first token time
|
||||
if firstTokenMs == nil && startsClientOutput {
|
||||
if !guardFirstOutput && firstTokenMs == nil && startsClientOutput {
|
||||
ms := int(time.Since(startTime).Milliseconds())
|
||||
firstTokenMs = &ms
|
||||
stopFirstOutputTimer()
|
||||
}
|
||||
s.parseSSEUsageBytes(dataBytes, usage)
|
||||
return
|
||||
}
|
||||
|
||||
// Forward non-data lines as-is. Flush only after the blank line that
|
||||
// completes the SSE event, never after a partial data line.
|
||||
// A blank line dispatches a guarded event from the attempt-local stage.
|
||||
if guardFirstOutput && line == "" {
|
||||
if !clientDisconnected {
|
||||
if _, err := writePendingString("\n"); err != nil {
|
||||
handlePendingWriteError(err)
|
||||
}
|
||||
}
|
||||
if streamEarlyErr == nil {
|
||||
completeGuardedEvent(queueDrained)
|
||||
}
|
||||
return
|
||||
}
|
||||
// Non-guarded streams retain upstream's event-boundary flushing: a keepalive
|
||||
// or queue-drain flush must never split an open SSE event.
|
||||
shouldFlush := false
|
||||
if line == "" {
|
||||
shouldFlush = eventShouldFlush || (queueDrained && clientOutputStarted)
|
||||
eventShouldFlush = false
|
||||
}
|
||||
if !clientDisconnected {
|
||||
if _, err := bufferedWriter.WriteString(line); err != nil {
|
||||
clientDisconnected = true
|
||||
logger.LegacyPrintf("service.openai_gateway", "Client disconnected during streaming, continuing to drain upstream for billing")
|
||||
} else if _, err := bufferedWriter.WriteString("\n"); err != nil {
|
||||
clientDisconnected = true
|
||||
logger.LegacyPrintf("service.openai_gateway", "Client disconnected during streaming, continuing to drain upstream for billing")
|
||||
if _, err := writePendingString(line); err != nil {
|
||||
handlePendingWriteError(err)
|
||||
} else if _, err := writePendingString("\n"); err != nil {
|
||||
handlePendingWriteError(err)
|
||||
} else {
|
||||
eventInProgress = line != ""
|
||||
if shouldFlush {
|
||||
@@ -404,7 +572,7 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
}
|
||||
|
||||
// 无超时/无 keepalive 的常见路径走同步扫描,减少 goroutine 与 channel 开销。
|
||||
if streamInterval <= 0 && keepaliveInterval <= 0 {
|
||||
if streamInterval <= 0 && keepaliveInterval <= 0 && firstOutputTimeout <= 0 {
|
||||
defer putSSEScannerBuf64K(scanBuf)
|
||||
for documentScanner.Scan() {
|
||||
processSSELine(documentScanner.Text(), true)
|
||||
@@ -419,20 +587,40 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
}
|
||||
|
||||
type scanEvent struct {
|
||||
line string
|
||||
err error
|
||||
line string
|
||||
err error
|
||||
processed chan struct{}
|
||||
}
|
||||
// 独立 goroutine 读取上游,避免读取阻塞影响 keepalive/超时处理
|
||||
events := make(chan scanEvent, 16)
|
||||
// Guard mode permits one queued token plus the token being processed. With
|
||||
// the guarded scanner cap this bounds scanner/channel retention near 16 MiB;
|
||||
// the timeout-disabled path preserves the legacy depth of 16.
|
||||
events := make(chan scanEvent, openAIFirstOutputEventQueueSize(guardFirstOutput))
|
||||
done := make(chan struct{})
|
||||
sendEvent := func(ev scanEvent) bool {
|
||||
if guardFirstOutput {
|
||||
ev.processed = make(chan struct{})
|
||||
}
|
||||
select {
|
||||
case events <- ev:
|
||||
case <-done:
|
||||
return false
|
||||
}
|
||||
if ev.processed == nil {
|
||||
return true
|
||||
}
|
||||
select {
|
||||
case <-ev.processed:
|
||||
return true
|
||||
case <-done:
|
||||
return false
|
||||
}
|
||||
}
|
||||
markEventProcessed := func(ev scanEvent) {
|
||||
if ev.processed != nil {
|
||||
close(ev.processed)
|
||||
}
|
||||
}
|
||||
var lastReadAt int64
|
||||
atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
|
||||
go func(scanBuf *sseScannerBuf64K) {
|
||||
@@ -454,12 +642,19 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
select {
|
||||
case ev, ok := <-events:
|
||||
if !ok {
|
||||
if guardFirstOutput && eventInProgress {
|
||||
// EOF dispatches the final SSE event even without a trailing blank
|
||||
// line. Do not synthesize extra bytes on the downstream wire.
|
||||
completeGuardedEvent(true)
|
||||
}
|
||||
return finalizeStream()
|
||||
}
|
||||
if result, err, done := handleScanErr(ev.err); done {
|
||||
markEventProcessed(ev)
|
||||
return result, err
|
||||
}
|
||||
processSSELine(ev.line, len(events) == 0)
|
||||
markEventProcessed(ev)
|
||||
if streamEarlyErr != nil {
|
||||
return resultWithUsage(), streamEarlyErr
|
||||
}
|
||||
@@ -480,6 +675,20 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
sendErrorEvent("stream_timeout")
|
||||
return resultWithUsage(), fmt.Errorf("stream data interval timeout")
|
||||
|
||||
case <-firstOutputCh:
|
||||
if firstTokenMs != nil {
|
||||
stopFirstOutputTimer()
|
||||
continue
|
||||
}
|
||||
_ = resp.Body.Close()
|
||||
for ev := range events {
|
||||
markEventProcessed(ev)
|
||||
}
|
||||
return resultWithUsage(), s.newOpenAIFirstOutputTimeoutError(
|
||||
ctx, c, account, startTime, originalModel, reasoningEffort,
|
||||
firstOutputTimeout, "semantic_output", resp.Header,
|
||||
)
|
||||
|
||||
case <-keepaliveCh:
|
||||
if clientDisconnected {
|
||||
continue
|
||||
@@ -490,7 +699,19 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
if time.Since(lastDownstreamWriteAt) < keepaliveInterval {
|
||||
continue
|
||||
}
|
||||
if _, err := bufferedWriter.WriteString(":\n\n"); err != nil {
|
||||
if guardFirstOutput {
|
||||
// Bypass attempt-local buffered frames. The stable SSE headers may be
|
||||
// committed here, but account headers remain private until semantic output.
|
||||
if _, err := w.Write([]byte(":\n\n")); err != nil {
|
||||
clientDisconnected = true
|
||||
logger.LegacyPrintf("service.openai_gateway", "Client disconnected during streaming, continuing to drain upstream for billing")
|
||||
continue
|
||||
}
|
||||
flusher.Flush()
|
||||
lastDownstreamWriteAt = time.Now()
|
||||
continue
|
||||
}
|
||||
if _, err := writePendingString(":\n\n"); err != nil {
|
||||
clientDisconnected = true
|
||||
logger.LegacyPrintf("service.openai_gateway", "Client disconnected during streaming, continuing to drain upstream for billing")
|
||||
continue
|
||||
|
||||
@@ -155,6 +155,15 @@ gateway:
|
||||
# OpenAI/Codex upstream response header timeout (seconds, 0=disabled)
|
||||
# OpenAI/Codex 等待上游响应头超时时间(秒,0=禁用本地响应头超时)
|
||||
openai_response_header_timeout: 0
|
||||
# Native OpenAI HTTP Responses first semantic output timeout (seconds, 0=disabled)
|
||||
# Includes response-header wait; does not apply to passthrough or WebSocket transports.
|
||||
# A timed-out request may already have incurred upstream usage; account failover can therefore duplicate upstream billing.
|
||||
# 超时请求可能已产生上游用量;切换账号重试可能导致上游重复计费。
|
||||
# Pre-output attempt staging is capped at 8 MiB; overflow fails over without exposing partial SSE data.
|
||||
# 首次输出前的单次尝试暂存上限为 8 MiB;溢出时切号且不暴露不完整 SSE 数据。
|
||||
openai_first_output_timeout_seconds: 0
|
||||
# Optional high/xhigh/max override (seconds, 0=use the standard timeout)
|
||||
openai_high_effort_first_output_timeout_seconds: 0
|
||||
# Max request body size in bytes (default: 256MB)
|
||||
# 请求体最大字节数(默认 256MB)
|
||||
max_body_size: 268435456
|
||||
|
||||
Reference in New Issue
Block a user