mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-11-03 00:33:13 +00:00
862 lines
31 KiB
Go
862 lines
31 KiB
Go
package service
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"net/http"
|
||
"net/url"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/tidwall/gjson"
|
||
)
|
||
|
||
func (s *OpenAIGatewayService) forwardOpenAIWSV2(
|
||
ctx context.Context,
|
||
c *gin.Context,
|
||
account *Account,
|
||
reqBody map[string]any,
|
||
clientPromptCacheKey string,
|
||
token string,
|
||
decision OpenAIWSProtocolDecision,
|
||
isCodexCLI bool,
|
||
reqStream bool,
|
||
originalModel string,
|
||
mappedModel string,
|
||
startTime time.Time,
|
||
attempt int,
|
||
lastFailureReason string,
|
||
agentTaskRecoveryTried *bool,
|
||
) (*OpenAIForwardResult, error) {
|
||
if s == nil || account == nil {
|
||
return nil, wrapOpenAIWSFallback("invalid_state", errors.New("service or account is nil"))
|
||
}
|
||
responseModelObserver := &upstreamResponseModelObserver{}
|
||
|
||
wsURL, err := s.buildOpenAIResponsesWSURL(account)
|
||
if err != nil {
|
||
return nil, wrapOpenAIWSFallback("build_ws_url", err)
|
||
}
|
||
wsHost := "-"
|
||
wsPath := "-"
|
||
if parsed, parseErr := url.Parse(wsURL); parseErr == nil && parsed != nil {
|
||
if h := strings.TrimSpace(parsed.Host); h != "" {
|
||
wsHost = normalizeOpenAIWSLogValue(h)
|
||
}
|
||
if p := strings.TrimSpace(parsed.Path); p != "" {
|
||
wsPath = normalizeOpenAIWSLogValue(p)
|
||
}
|
||
}
|
||
logOpenAIWSModeDebug(
|
||
"dial_target account_id=%d account_type=%s ws_host=%s ws_path=%s",
|
||
account.ID,
|
||
account.Type,
|
||
wsHost,
|
||
wsPath,
|
||
)
|
||
|
||
payload := s.buildOpenAIWSCreatePayload(reqBody, account)
|
||
payloadStrategy, removedKeys := applyOpenAIWSRetryPayloadStrategy(payload, attempt)
|
||
turnState := ""
|
||
turnMetadata := ""
|
||
if c != nil && c.Request != nil {
|
||
turnState = strings.TrimSpace(c.GetHeader(openAIWSTurnStateHeader))
|
||
turnMetadata = strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader))
|
||
}
|
||
setOpenAIWSTurnMetadata(payload, turnMetadata)
|
||
applyStagedCodexFingerprintClientMetadata(c, account, payload)
|
||
previousResponseID := openAIWSPayloadString(payload, "previous_response_id")
|
||
previousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(previousResponseID)
|
||
promptCacheKey := strings.TrimSpace(clientPromptCacheKey)
|
||
if promptCacheKey == "" {
|
||
// Fingerprint convergence may inject a default key when the client did
|
||
// not send one; retain that fallback without replacing an explicit raw key.
|
||
promptCacheKey = openAIWSPayloadString(payload, "prompt_cache_key")
|
||
}
|
||
_, hasTools := payload["tools"]
|
||
debugEnabled := isOpenAIWSModeDebugEnabled()
|
||
payloadBytes := -1
|
||
resolvePayloadBytes := func() int {
|
||
if payloadBytes >= 0 {
|
||
return payloadBytes
|
||
}
|
||
payloadBytes = len(payloadAsJSONBytes(payload))
|
||
return payloadBytes
|
||
}
|
||
streamValue := "-"
|
||
if raw, ok := payload["stream"]; ok {
|
||
streamValue = normalizeOpenAIWSLogValue(strings.TrimSpace(fmt.Sprintf("%v", raw)))
|
||
}
|
||
payloadEventType := openAIWSPayloadString(payload, "type")
|
||
if payloadEventType == "" {
|
||
payloadEventType = "response.create"
|
||
}
|
||
if s.shouldEmitOpenAIWSPayloadSchema(attempt) {
|
||
logOpenAIWSModeInfo(
|
||
"[debug] payload_schema account_id=%d attempt=%d event=%s payload_keys=%s payload_bytes=%d payload_key_sizes=%s input_summary=%s stream=%s payload_strategy=%s removed_keys=%s has_previous_response_id=%v has_prompt_cache_key=%v has_tools=%v",
|
||
account.ID,
|
||
attempt,
|
||
payloadEventType,
|
||
normalizeOpenAIWSLogValue(strings.Join(sortedKeys(payload), ",")),
|
||
resolvePayloadBytes(),
|
||
normalizeOpenAIWSLogValue(summarizeOpenAIWSPayloadKeySizes(payload, openAIWSPayloadKeySizeTopN)),
|
||
normalizeOpenAIWSLogValue(summarizeOpenAIWSInput(payload["input"])),
|
||
streamValue,
|
||
normalizeOpenAIWSLogValue(payloadStrategy),
|
||
normalizeOpenAIWSLogValue(strings.Join(removedKeys, ",")),
|
||
previousResponseID != "",
|
||
promptCacheKey != "",
|
||
hasTools,
|
||
)
|
||
}
|
||
|
||
stateStore := s.getOpenAIWSStateStore()
|
||
groupID := getOpenAIGroupIDFromContext(c)
|
||
sessionHash := s.GenerateSessionHash(c, nil)
|
||
if sessionHash == "" {
|
||
var legacySessionHash string
|
||
sessionHash, legacySessionHash = openAIWSSessionHashesFromID(promptCacheKey)
|
||
attachOpenAILegacySessionHashToGin(c, legacySessionHash)
|
||
}
|
||
if turnState == "" && stateStore != nil && sessionHash != "" {
|
||
if savedTurnState, ok := stateStore.GetSessionTurnState(groupID, sessionHash); ok {
|
||
turnState = savedTurnState
|
||
}
|
||
}
|
||
preferredConnID := ""
|
||
if stateStore != nil && previousResponseID != "" {
|
||
if connID, ok := stateStore.GetResponseConn(previousResponseID); ok {
|
||
preferredConnID = connID
|
||
}
|
||
}
|
||
storeDisabled := s.isOpenAIWSStoreDisabledInRequest(reqBody, account)
|
||
if stateStore != nil && storeDisabled && previousResponseID == "" && sessionHash != "" {
|
||
if connID, ok := stateStore.GetSessionConn(groupID, sessionHash); ok {
|
||
preferredConnID = connID
|
||
}
|
||
}
|
||
storeDisabledConnMode := s.openAIWSStoreDisabledConnMode()
|
||
forceNewConnByPolicy := shouldForceNewConnOnStoreDisabled(storeDisabledConnMode, lastFailureReason)
|
||
forceNewConn := forceNewConnByPolicy && storeDisabled && previousResponseID == "" && sessionHash != "" && preferredConnID == ""
|
||
wsHeaders, sessionResolution, buildHdrErr := s.buildOpenAIWSHeaders(
|
||
ctx,
|
||
c,
|
||
account,
|
||
token,
|
||
decision,
|
||
isCodexCLI,
|
||
turnState,
|
||
turnMetadata,
|
||
promptCacheKey,
|
||
openAIWSPayloadString(payload, "model"),
|
||
openAIWSPayloadString(payload, "service_tier"),
|
||
)
|
||
if buildHdrErr != nil {
|
||
return nil, fmt.Errorf("build ws headers: %w", buildHdrErr)
|
||
}
|
||
logOpenAIWSModeDebug(
|
||
"acquire_start account_id=%d account_type=%s transport=%s preferred_conn_id=%s has_previous_response_id=%v session_hash=%s has_turn_state=%v turn_state_len=%d has_turn_metadata=%v turn_metadata_len=%d store_disabled=%v store_disabled_conn_mode=%s retry_last_reason=%s force_new_conn=%v header_user_agent=%s header_openai_beta=%s header_originator=%s header_accept_language=%s header_session_id=%s header_conversation_id=%s session_id_source=%s conversation_id_source=%s has_prompt_cache_key=%v has_chatgpt_account_id=%v has_authorization=%v has_session_id=%v has_conversation_id=%v proxy_enabled=%v",
|
||
account.ID,
|
||
account.Type,
|
||
normalizeOpenAIWSLogValue(string(decision.Transport)),
|
||
truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen),
|
||
previousResponseID != "",
|
||
truncateOpenAIWSLogValue(sessionHash, 12),
|
||
turnState != "",
|
||
len(turnState),
|
||
turnMetadata != "",
|
||
len(turnMetadata),
|
||
storeDisabled,
|
||
normalizeOpenAIWSLogValue(storeDisabledConnMode),
|
||
truncateOpenAIWSLogValue(lastFailureReason, openAIWSLogValueMaxLen),
|
||
forceNewConn,
|
||
openAIWSHeaderValueForLog(wsHeaders, "user-agent"),
|
||
openAIWSHeaderValueForLog(wsHeaders, "openai-beta"),
|
||
openAIWSHeaderValueForLog(wsHeaders, "originator"),
|
||
openAIWSHeaderValueForLog(wsHeaders, "accept-language"),
|
||
openAIWSHeaderValueForLog(wsHeaders, "session_id"),
|
||
openAIWSHeaderValueForLog(wsHeaders, "conversation_id"),
|
||
normalizeOpenAIWSLogValue(sessionResolution.SessionSource),
|
||
normalizeOpenAIWSLogValue(sessionResolution.ConversationSource),
|
||
promptCacheKey != "",
|
||
hasOpenAIWSHeader(wsHeaders, "chatgpt-account-id"),
|
||
hasOpenAIWSHeader(wsHeaders, "authorization"),
|
||
hasOpenAIWSHeader(wsHeaders, "session_id"),
|
||
hasOpenAIWSHeader(wsHeaders, "conversation_id"),
|
||
account.ProxyID != nil && account.Proxy != nil,
|
||
)
|
||
|
||
acquireCtx, acquireCancel := context.WithTimeout(ctx, s.openAIWSAcquireTimeout())
|
||
defer acquireCancel()
|
||
|
||
lease, err := s.getOpenAIWSConnPool().Acquire(acquireCtx, openAIWSAcquireRequest{
|
||
Account: account,
|
||
WSURL: wsURL,
|
||
Headers: wsHeaders,
|
||
HeadersFactory: func(factoryCtx context.Context, headers http.Header) (http.Header, error) {
|
||
return s.refreshOpenAIAgentIdentityHeaders(factoryCtx, account, headers)
|
||
},
|
||
PreferredConnID: preferredConnID,
|
||
ForceNewConn: forceNewConn,
|
||
ProxyURL: func() string {
|
||
if account.ProxyID != nil && account.Proxy != nil {
|
||
return account.Proxy.URL()
|
||
}
|
||
return ""
|
||
}(),
|
||
})
|
||
if err != nil {
|
||
var agentDialErr *openAIWSDialError
|
||
if s.isAgentIdentityAccount(ctx, account) && errors.As(err, &agentDialErr) && isAgentIdentityTaskInvalidWSDialError(agentDialErr) && agentTaskRecoveryTried != nil && !*agentTaskRecoveryTried {
|
||
*agentTaskRecoveryTried = true
|
||
if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil {
|
||
return nil, fmt.Errorf("agent identity task recovery failed: %w", recoveryErr)
|
||
}
|
||
return nil, &agentIdentityTaskRecoveredError{}
|
||
}
|
||
s.handleOpenAIWSDialTransientFailure(ctx, account, mappedModel, err)
|
||
dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(err)
|
||
logOpenAIWSModeInfo(
|
||
"acquire_fail account_id=%d account_type=%s transport=%s reason=%s dial_status=%d dial_class=%s dial_close_status=%s dial_close_reason=%s dial_resp_server=%s dial_resp_via=%s dial_resp_cf_ray=%s dial_resp_x_request_id=%s cause=%s preferred_conn_id=%s force_new_conn=%v ws_host=%s ws_path=%s proxy_enabled=%v",
|
||
account.ID,
|
||
account.Type,
|
||
normalizeOpenAIWSLogValue(string(decision.Transport)),
|
||
normalizeOpenAIWSLogValue(classifyOpenAIWSAcquireError(err)),
|
||
dialStatus,
|
||
dialClass,
|
||
dialCloseStatus,
|
||
truncateOpenAIWSLogValue(dialCloseReason, openAIWSHeaderValueMaxLen),
|
||
dialRespServer,
|
||
dialRespVia,
|
||
dialRespCFRay,
|
||
dialRespReqID,
|
||
truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen),
|
||
truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen),
|
||
forceNewConn,
|
||
wsHost,
|
||
wsPath,
|
||
account.ProxyID != nil && account.Proxy != nil,
|
||
)
|
||
var dialErr *openAIWSDialError
|
||
if errors.As(err, &dialErr) && dialErr != nil && dialErr.StatusCode == http.StatusTooManyRequests {
|
||
s.persistOpenAIWSRateLimitSignal(ctx, account, dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(err.Error()), mappedModel)
|
||
}
|
||
return nil, wrapOpenAIWSFallback(classifyOpenAIWSAcquireError(err), err)
|
||
}
|
||
// cleanExit 标记正常终端事件退出，此时上游不会再发送帧，连接可安全归还复用。
|
||
// 所有异常路径（读写错误、error 事件等）已在各自分支中提前调用 MarkBroken，
|
||
// 因此 defer 中只需处理正常退出时不 MarkBroken 即可。
|
||
cleanExit := false
|
||
defer func() {
|
||
if !cleanExit {
|
||
lease.MarkBroken()
|
||
}
|
||
lease.Release()
|
||
}()
|
||
connID := strings.TrimSpace(lease.ConnID())
|
||
logOpenAIWSModeDebug(
|
||
"connected account_id=%d account_type=%s transport=%s conn_id=%s conn_reused=%v conn_pick_ms=%d queue_wait_ms=%d has_previous_response_id=%v",
|
||
account.ID,
|
||
account.Type,
|
||
normalizeOpenAIWSLogValue(string(decision.Transport)),
|
||
connID,
|
||
lease.Reused(),
|
||
lease.ConnPickDuration().Milliseconds(),
|
||
lease.QueueWaitDuration().Milliseconds(),
|
||
previousResponseID != "",
|
||
)
|
||
if previousResponseID != "" {
|
||
logOpenAIWSModeInfo(
|
||
"continuation_probe account_id=%d account_type=%s conn_id=%s previous_response_id=%s previous_response_id_kind=%s preferred_conn_id=%s conn_reused=%v store_disabled=%v session_hash=%s header_session_id=%s header_conversation_id=%s session_id_source=%s conversation_id_source=%s has_turn_state=%v turn_state_len=%d has_prompt_cache_key=%v",
|
||
account.ID,
|
||
account.Type,
|
||
truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen),
|
||
truncateOpenAIWSLogValue(previousResponseID, openAIWSIDValueMaxLen),
|
||
normalizeOpenAIWSLogValue(previousResponseIDKind),
|
||
truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen),
|
||
lease.Reused(),
|
||
storeDisabled,
|
||
truncateOpenAIWSLogValue(sessionHash, 12),
|
||
openAIWSHeaderValueForLog(wsHeaders, "session_id"),
|
||
openAIWSHeaderValueForLog(wsHeaders, "conversation_id"),
|
||
normalizeOpenAIWSLogValue(sessionResolution.SessionSource),
|
||
normalizeOpenAIWSLogValue(sessionResolution.ConversationSource),
|
||
turnState != "",
|
||
len(turnState),
|
||
promptCacheKey != "",
|
||
)
|
||
}
|
||
if c != nil {
|
||
SetOpsLatencyMs(c, OpsOpenAIWSConnPickMsKey, lease.ConnPickDuration().Milliseconds())
|
||
SetOpsLatencyMs(c, OpsOpenAIWSQueueWaitMsKey, lease.QueueWaitDuration().Milliseconds())
|
||
c.Set(OpsOpenAIWSConnReusedKey, lease.Reused())
|
||
if connID != "" {
|
||
c.Set(OpsOpenAIWSConnIDKey, connID)
|
||
}
|
||
}
|
||
|
||
handshakeTurnState := strings.TrimSpace(lease.HandshakeHeader(openAIWSTurnStateHeader))
|
||
logOpenAIWSModeDebug(
|
||
"handshake account_id=%d conn_id=%s has_turn_state=%v turn_state_len=%d",
|
||
account.ID,
|
||
connID,
|
||
handshakeTurnState != "",
|
||
len(handshakeTurnState),
|
||
)
|
||
if handshakeTurnState != "" {
|
||
if stateStore != nil && sessionHash != "" {
|
||
stateStore.BindSessionTurnState(groupID, sessionHash, handshakeTurnState, s.openAIWSSessionStickyTTL())
|
||
}
|
||
if c != nil {
|
||
c.Header(http.CanonicalHeaderKey(openAIWSTurnStateHeader), handshakeTurnState)
|
||
}
|
||
}
|
||
|
||
if err := s.performOpenAIWSGeneratePrewarm(
|
||
ctx,
|
||
lease,
|
||
decision,
|
||
payload,
|
||
previousResponseID,
|
||
reqBody,
|
||
account,
|
||
stateStore,
|
||
groupID,
|
||
); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
if err := lease.WriteJSONWithContextTimeout(ctx, payload, s.openAIWSWriteTimeout()); err != nil {
|
||
lease.MarkBroken()
|
||
logOpenAIWSModeInfo(
|
||
"write_request_fail account_id=%d conn_id=%s cause=%s payload_bytes=%d",
|
||
account.ID,
|
||
connID,
|
||
truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen),
|
||
resolvePayloadBytes(),
|
||
)
|
||
return nil, wrapOpenAIWSFallback("write_request", err)
|
||
}
|
||
if debugEnabled {
|
||
logOpenAIWSModeDebug(
|
||
"write_request_sent account_id=%d conn_id=%s stream=%v payload_bytes=%d previous_response_id=%s",
|
||
account.ID,
|
||
connID,
|
||
reqStream,
|
||
resolvePayloadBytes(),
|
||
truncateOpenAIWSLogValue(previousResponseID, openAIWSIDValueMaxLen),
|
||
)
|
||
}
|
||
|
||
usage := &OpenAIUsage{}
|
||
imageCounter := newOpenAIImageOutputCounter()
|
||
var firstTokenMs *int
|
||
responseID := ""
|
||
var finalResponse []byte
|
||
wroteDownstream := false
|
||
needModelReplace := originalModel != mappedModel
|
||
var mappedModelBytes []byte
|
||
if needModelReplace && mappedModel != "" {
|
||
mappedModelBytes = []byte(mappedModel)
|
||
}
|
||
bufferedStreamEvents := make([][]byte, 0, 4)
|
||
eventCount := 0
|
||
tokenEventCount := 0
|
||
terminalEventCount := 0
|
||
bufferedEventCount := 0
|
||
flushedBufferedEventCount := 0
|
||
firstEventType := ""
|
||
lastEventType := ""
|
||
upstreamTerminalEvent := ""
|
||
clientDisconnected := false
|
||
clientDisconnectDrainStartedAt := time.Time{}
|
||
readTimeout := s.openAIWSReadTimeout()
|
||
upstreamReadCtx := ctx
|
||
upstreamReadDetached := false
|
||
clientRequestCanceled := func() bool {
|
||
return ctx != nil && errors.Is(ctx.Err(), context.Canceled)
|
||
}
|
||
markClientDisconnected := func(cause string) {
|
||
if clientDisconnected {
|
||
return
|
||
}
|
||
clientDisconnected = true
|
||
clientDisconnectDrainStartedAt = time.Now()
|
||
if !upstreamReadDetached {
|
||
upstreamReadCtx = context.WithoutCancel(ctx)
|
||
upstreamReadDetached = true
|
||
}
|
||
logOpenAIWSModeInfo(
|
||
"client_disconnected account_id=%d conn_id=%s cause=%s events=%d token_events=%d",
|
||
account.ID,
|
||
connID,
|
||
cause,
|
||
eventCount,
|
||
tokenEventCount,
|
||
)
|
||
}
|
||
markClientRequestCanceled := func() {
|
||
if clientRequestCanceled() {
|
||
markClientDisconnected("request_context_canceled")
|
||
}
|
||
}
|
||
resultWithUsage := func() *OpenAIForwardResult {
|
||
return &OpenAIForwardResult{
|
||
RequestID: responseID,
|
||
ResponseID: responseID,
|
||
Usage: *usage,
|
||
Model: originalModel,
|
||
UpstreamModel: mappedModel,
|
||
UpstreamResponseModel: responseModelObserver.Model(),
|
||
UpstreamResponseModelConflict: responseModelObserver.Conflict(),
|
||
UpstreamResponseServiceTier: responseModelObserver.ServiceTier(),
|
||
ServiceTier: resolvedOpenAIUpstreamServiceTierFromObserver(responseModelObserver, extractOpenAIServiceTier(reqBody)),
|
||
ReasoningEffort: extractOpenAIReasoningEffort(reqBody, mappedModel, originalModel),
|
||
Stream: reqStream,
|
||
OpenAIWSMode: true,
|
||
UpstreamTerminalEvent: upstreamTerminalEvent,
|
||
ResponseHeaders: lease.HandshakeHeaders(),
|
||
Duration: time.Since(startTime),
|
||
FirstTokenMs: firstTokenMs,
|
||
ClientDisconnect: clientDisconnected,
|
||
}
|
||
}
|
||
|
||
var flusher http.Flusher
|
||
if reqStream {
|
||
if s.responseHeaderFilter != nil {
|
||
responseheaders.WriteFilteredHeaders(c.Writer.Header(), http.Header{}, s.responseHeaderFilter)
|
||
}
|
||
c.Header("Content-Type", "text/event-stream")
|
||
c.Header("Cache-Control", "no-cache")
|
||
c.Header("Connection", "keep-alive")
|
||
c.Header("X-Accel-Buffering", "no")
|
||
f, ok := c.Writer.(http.Flusher)
|
||
if !ok {
|
||
lease.MarkBroken()
|
||
return nil, wrapOpenAIWSFallback("streaming_not_supported", errors.New("streaming not supported"))
|
||
}
|
||
flusher = f
|
||
}
|
||
|
||
flushBatchSize := s.openAIWSEventFlushBatchSize()
|
||
flushInterval := s.openAIWSEventFlushInterval()
|
||
pendingFlushEvents := 0
|
||
lastFlushAt := time.Now()
|
||
flushStreamWriter := func(force bool) {
|
||
if clientDisconnected || flusher == nil || pendingFlushEvents <= 0 {
|
||
return
|
||
}
|
||
if !force && flushBatchSize > 1 && pendingFlushEvents < flushBatchSize {
|
||
if flushInterval <= 0 || time.Since(lastFlushAt) < flushInterval {
|
||
return
|
||
}
|
||
}
|
||
flusher.Flush()
|
||
pendingFlushEvents = 0
|
||
lastFlushAt = time.Now()
|
||
}
|
||
emitStreamMessage := func(message []byte, forceFlush bool) {
|
||
if clientDisconnected {
|
||
return
|
||
}
|
||
frame := make([]byte, 0, len(message)+8)
|
||
frame = append(frame, "data: "...)
|
||
frame = append(frame, message...)
|
||
frame = append(frame, '\n', '\n')
|
||
_, wErr := c.Writer.Write(frame)
|
||
if wErr == nil {
|
||
wroteDownstream = true
|
||
pendingFlushEvents++
|
||
flushStreamWriter(forceFlush)
|
||
return
|
||
}
|
||
markClientDisconnected("downstream_write_error")
|
||
logger.LegacyPrintf("service.openai_gateway", "[OpenAI WS Mode] client disconnected, continue draining upstream: account=%d", account.ID)
|
||
}
|
||
flushBufferedStreamEvents := func(reason string) {
|
||
if len(bufferedStreamEvents) == 0 {
|
||
return
|
||
}
|
||
flushed := len(bufferedStreamEvents)
|
||
for _, buffered := range bufferedStreamEvents {
|
||
emitStreamMessage(buffered, false)
|
||
}
|
||
bufferedStreamEvents = bufferedStreamEvents[:0]
|
||
flushStreamWriter(true)
|
||
flushedBufferedEventCount += flushed
|
||
if debugEnabled {
|
||
logOpenAIWSModeDebug(
|
||
"buffer_flush account_id=%d conn_id=%s reason=%s flushed=%d total_flushed=%d client_disconnected=%v",
|
||
account.ID,
|
||
connID,
|
||
truncateOpenAIWSLogValue(reason, openAIWSLogValueMaxLen),
|
||
flushed,
|
||
flushedBufferedEventCount,
|
||
clientDisconnected,
|
||
)
|
||
}
|
||
}
|
||
|
||
// Keep per-read timeouts unchanged for connected clients. Once a client
|
||
// disconnects, use the same timeout as a bounded total drain budget.
|
||
var pendingJSONDocuments [][]byte
|
||
|
||
readLoop:
|
||
for {
|
||
markClientRequestCanceled()
|
||
var message []byte
|
||
var readErr error
|
||
readUsedDetachedContext := upstreamReadDetached
|
||
if len(pendingJSONDocuments) > 0 {
|
||
message = pendingJSONDocuments[0]
|
||
pendingJSONDocuments = pendingJSONDocuments[1:]
|
||
} else {
|
||
currentReadTimeout := readTimeout
|
||
if clientDisconnected && !clientDisconnectDrainStartedAt.IsZero() {
|
||
remaining := readTimeout - time.Since(clientDisconnectDrainStartedAt)
|
||
if remaining <= 0 {
|
||
lease.MarkBroken()
|
||
break readLoop
|
||
}
|
||
if remaining < currentReadTimeout {
|
||
currentReadTimeout = remaining
|
||
}
|
||
}
|
||
message, readErr = lease.ReadMessageWithContextTimeout(upstreamReadCtx, currentReadTimeout)
|
||
if readErr == nil {
|
||
if documents, repaired := splitOpenAIConcatenatedJSONDocuments(message); repaired {
|
||
logOpenAIWSModeInfo(
|
||
"concatenated_json_repaired account_id=%d conn_id=%s documents=%d bytes=%d",
|
||
account.ID,
|
||
truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen),
|
||
len(documents),
|
||
len(message),
|
||
)
|
||
message = documents[0]
|
||
pendingJSONDocuments = append(pendingJSONDocuments, documents[1:]...)
|
||
}
|
||
}
|
||
}
|
||
markClientRequestCanceled()
|
||
if readErr == nil && !json.Valid(message) {
|
||
eventType, _, _ := parseOpenAIWSEventEnvelope(message)
|
||
if eventType == "" {
|
||
eventType = "unknown"
|
||
}
|
||
lease.MarkBroken()
|
||
logOpenAIWSModeInfo(
|
||
"invalid_event_json account_id=%d conn_id=%s event_type=%s bytes=%d wrote_downstream=%v",
|
||
account.ID,
|
||
truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen),
|
||
truncateOpenAIWSLogValue(eventType, openAIWSLogValueMaxLen),
|
||
len(message),
|
||
wroteDownstream,
|
||
)
|
||
if !wroteDownstream {
|
||
return nil, wrapOpenAIWSFallback("invalid_event_json", errors.New("upstream websocket returned malformed Responses event JSON"))
|
||
}
|
||
return nil, errors.New("upstream websocket returned malformed Responses event JSON after downstream output")
|
||
}
|
||
if readErr != nil {
|
||
lease.MarkBroken()
|
||
closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr)
|
||
logOpenAIWSModeInfo(
|
||
"read_fail account_id=%d conn_id=%s wrote_downstream=%v close_status=%s close_reason=%s cause=%s events=%d token_events=%d terminal_events=%d buffered_pending=%d buffered_flushed=%d first_event=%s last_event=%s",
|
||
account.ID,
|
||
connID,
|
||
wroteDownstream,
|
||
closeStatus,
|
||
closeReason,
|
||
truncateOpenAIWSLogValue(readErr.Error(), openAIWSLogValueMaxLen),
|
||
eventCount,
|
||
tokenEventCount,
|
||
terminalEventCount,
|
||
len(bufferedStreamEvents),
|
||
flushedBufferedEventCount,
|
||
truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen),
|
||
truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen),
|
||
)
|
||
if clientDisconnected {
|
||
if !readUsedDetachedContext && errors.Is(readErr, context.Canceled) && clientRequestCanceled() {
|
||
continue
|
||
}
|
||
break
|
||
}
|
||
if !wroteDownstream {
|
||
return nil, wrapOpenAIWSFallback(classifyOpenAIWSReadFallbackReason(readErr), readErr)
|
||
}
|
||
setOpsUpstreamError(c, 0, sanitizeUpstreamErrorMessage(readErr.Error()), "")
|
||
return nil, fmt.Errorf("openai ws read event: %w", readErr)
|
||
}
|
||
if normalized, changed := normalizeCompletedImageGenerationStatus(message); changed {
|
||
message = normalized
|
||
}
|
||
|
||
eventType, eventResponseID, responseField := parseOpenAIWSEventEnvelope(message)
|
||
if eventType == "" {
|
||
continue
|
||
}
|
||
responseModelObserver.ObserveOpenAI(message, eventType)
|
||
eventCount++
|
||
if firstEventType == "" {
|
||
firstEventType = eventType
|
||
}
|
||
lastEventType = eventType
|
||
|
||
if responseID == "" && eventResponseID != "" {
|
||
responseID = eventResponseID
|
||
}
|
||
|
||
isTokenEvent := isOpenAIWSTokenEvent(eventType)
|
||
if isTokenEvent {
|
||
tokenEventCount++
|
||
}
|
||
isTerminalEvent := isOpenAIWSTerminalEvent(eventType)
|
||
if isTerminalEvent {
|
||
terminalEventCount++
|
||
}
|
||
if firstTokenMs == nil && isTokenEvent {
|
||
ms := int(time.Since(startTime).Milliseconds())
|
||
firstTokenMs = &ms
|
||
}
|
||
if debugEnabled && shouldLogOpenAIWSEvent(eventCount, eventType) {
|
||
logOpenAIWSModeDebug(
|
||
"event_received account_id=%d conn_id=%s idx=%d type=%s bytes=%d token=%v terminal=%v buffered_pending=%d",
|
||
account.ID,
|
||
connID,
|
||
eventCount,
|
||
truncateOpenAIWSLogValue(eventType, openAIWSLogValueMaxLen),
|
||
len(message),
|
||
isTokenEvent,
|
||
isTerminalEvent,
|
||
len(bufferedStreamEvents),
|
||
)
|
||
}
|
||
|
||
if !clientDisconnected {
|
||
if needModelReplace && len(mappedModelBytes) > 0 && openAIWSEventMayContainModel(eventType) && bytes.Contains(message, mappedModelBytes) {
|
||
message = replaceOpenAIWSMessageModel(message, mappedModel, originalModel)
|
||
}
|
||
if openAIWSEventMayContainToolCalls(eventType) && openAIWSMessageLikelyContainsToolCalls(message) {
|
||
if corrected, changed := s.toolCorrector.CorrectToolCallsInSSEBytes(message); changed {
|
||
message = corrected
|
||
}
|
||
}
|
||
message = restoreCodexToolNamesFromContext(c, message)
|
||
}
|
||
if openAIWSMessageShouldParseUsage(eventType, message) {
|
||
parseOpenAIWSResponseUsageFromCompletedEvent(message, usage)
|
||
}
|
||
imageCounter.AddSSEData(message)
|
||
|
||
if eventType == "error" || eventType == "response.failed" {
|
||
markOpenAICyberPolicyEvent(c, message, http.StatusOK, usage)
|
||
}
|
||
|
||
if eventType == "error" {
|
||
s.handleOpenAIWSErrorEventTransientFailure(ctx, account, mappedModel, lease.HandshakeHeaders(), message)
|
||
errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(message)
|
||
s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw, mappedModel)
|
||
errMsg := strings.TrimSpace(errMsgRaw)
|
||
if errMsg == "" {
|
||
errMsg = "Upstream websocket error"
|
||
}
|
||
fallbackReason, canFallback := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMsgRaw)
|
||
errCode, errType, errMessage := summarizeOpenAIWSErrorEventFieldsFromRaw(errCodeRaw, errTypeRaw, errMsgRaw)
|
||
logOpenAIWSModeInfo(
|
||
"error_event account_id=%d conn_id=%s idx=%d fallback_reason=%s can_fallback=%v err_code=%s err_type=%s err_message=%s",
|
||
account.ID,
|
||
connID,
|
||
eventCount,
|
||
truncateOpenAIWSLogValue(fallbackReason, openAIWSLogValueMaxLen),
|
||
canFallback,
|
||
errCode,
|
||
errType,
|
||
errMessage,
|
||
)
|
||
if fallbackReason == "previous_response_not_found" {
|
||
logOpenAIWSModeInfo(
|
||
"previous_response_not_found_diag account_id=%d account_type=%s conn_id=%s previous_response_id=%s previous_response_id_kind=%s response_id=%s event_idx=%d req_stream=%v store_disabled=%v conn_reused=%v session_hash=%s header_session_id=%s header_conversation_id=%s session_id_source=%s conversation_id_source=%s has_turn_state=%v turn_state_len=%d has_prompt_cache_key=%v err_code=%s err_type=%s err_message=%s",
|
||
account.ID,
|
||
account.Type,
|
||
connID,
|
||
truncateOpenAIWSLogValue(previousResponseID, openAIWSIDValueMaxLen),
|
||
normalizeOpenAIWSLogValue(previousResponseIDKind),
|
||
truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen),
|
||
eventCount,
|
||
reqStream,
|
||
storeDisabled,
|
||
lease.Reused(),
|
||
truncateOpenAIWSLogValue(sessionHash, 12),
|
||
openAIWSHeaderValueForLog(wsHeaders, "session_id"),
|
||
openAIWSHeaderValueForLog(wsHeaders, "conversation_id"),
|
||
normalizeOpenAIWSLogValue(sessionResolution.SessionSource),
|
||
normalizeOpenAIWSLogValue(sessionResolution.ConversationSource),
|
||
turnState != "",
|
||
len(turnState),
|
||
promptCacheKey != "",
|
||
errCode,
|
||
errType,
|
||
errMessage,
|
||
)
|
||
}
|
||
// error 事件后连接不再可复用，避免回池后污染下一请求。
|
||
lease.MarkBroken()
|
||
if !wroteDownstream && canFallback {
|
||
return nil, wrapOpenAIWSFallback(fallbackReason, errors.New(errMsg))
|
||
}
|
||
statusCode := openAIWSErrorHTTPStatusFromRaw(errCodeRaw, errTypeRaw)
|
||
setOpsUpstreamError(c, statusCode, errMsg, "")
|
||
if reqStream && !clientDisconnected {
|
||
flushBufferedStreamEvents("error_event")
|
||
emitStreamMessage(message, true)
|
||
}
|
||
if !reqStream {
|
||
c.JSON(statusCode, gin.H{
|
||
"error": gin.H{
|
||
"type": "upstream_error",
|
||
"message": errMsg,
|
||
},
|
||
})
|
||
}
|
||
return nil, fmt.Errorf("openai ws error event: %s", errMsg)
|
||
}
|
||
|
||
if reqStream {
|
||
// 在首个 token 前先缓冲事件（如 response.created），
|
||
// 以便上游早期断连时仍可安全回退到 HTTP，不给下游发送半截流。
|
||
shouldBuffer := firstTokenMs == nil && !isTokenEvent && !isTerminalEvent
|
||
if shouldBuffer {
|
||
buffered := make([]byte, len(message))
|
||
copy(buffered, message)
|
||
bufferedStreamEvents = append(bufferedStreamEvents, buffered)
|
||
bufferedEventCount++
|
||
if debugEnabled && shouldLogOpenAIWSBufferedEvent(bufferedEventCount) {
|
||
logOpenAIWSModeDebug(
|
||
"buffer_enqueue account_id=%d conn_id=%s idx=%d event_idx=%d event_type=%s buffer_size=%d",
|
||
account.ID,
|
||
connID,
|
||
bufferedEventCount,
|
||
eventCount,
|
||
truncateOpenAIWSLogValue(eventType, openAIWSLogValueMaxLen),
|
||
len(bufferedStreamEvents),
|
||
)
|
||
}
|
||
} else {
|
||
flushBufferedStreamEvents(eventType)
|
||
emitStreamMessage(message, isTerminalEvent)
|
||
}
|
||
} else {
|
||
if responseField.Exists() && responseField.Type == gjson.JSON {
|
||
finalResponse = []byte(responseField.Raw)
|
||
}
|
||
}
|
||
|
||
if isTerminalEvent {
|
||
if !clientDisconnected {
|
||
markOpenAIWSClientVisibleFailure(c, eventType, message)
|
||
}
|
||
upstreamTerminalEvent = s.handleOpenAIWSTerminalTransientFailure(ctx, account, mappedModel, lease.HandshakeHeaders(), message)
|
||
// A terminal event must be the final JSON document in its WS message.
|
||
// Ignore any tail for the completed client turn, but never reuse the
|
||
// ambiguous upstream connection for another request.
|
||
cleanExit = len(pendingJSONDocuments) == 0
|
||
break
|
||
}
|
||
}
|
||
|
||
if clientDisconnected && terminalEventCount == 0 {
|
||
return resultWithUsage(), fmt.Errorf("openai ws stream incomplete after client disconnect: %w", context.Canceled)
|
||
}
|
||
if !reqStream {
|
||
if clientDisconnected {
|
||
return resultWithUsage(), nil
|
||
}
|
||
if len(finalResponse) == 0 {
|
||
logOpenAIWSModeInfo(
|
||
"missing_final_response account_id=%d conn_id=%s events=%d token_events=%d terminal_events=%d wrote_downstream=%v",
|
||
account.ID,
|
||
connID,
|
||
eventCount,
|
||
tokenEventCount,
|
||
terminalEventCount,
|
||
wroteDownstream,
|
||
)
|
||
if !wroteDownstream {
|
||
return nil, wrapOpenAIWSFallback("missing_final_response", errors.New("no terminal response payload"))
|
||
}
|
||
return nil, errors.New("ws finished without final response")
|
||
}
|
||
|
||
if needModelReplace {
|
||
finalResponse = s.replaceModelInResponseBody(finalResponse, mappedModel, originalModel)
|
||
}
|
||
finalResponse = s.correctToolCallsInResponseBody(finalResponse)
|
||
populateOpenAIUsageFromResponseJSON(finalResponse, usage)
|
||
if responseID == "" {
|
||
responseID = strings.TrimSpace(gjson.GetBytes(finalResponse, "id").String())
|
||
}
|
||
|
||
c.Data(http.StatusOK, "application/json", finalResponse)
|
||
} else {
|
||
flushStreamWriter(true)
|
||
}
|
||
|
||
if responseID != "" && stateStore != nil {
|
||
ttl := s.openAIWSResponseStickyTTL()
|
||
logOpenAIWSBindResponseAccountWarn(groupID, account.ID, responseID, stateStore.BindResponseAccount(ctx, groupID, responseID, account.ID, ttl))
|
||
stateStore.BindResponseConn(responseID, lease.ConnID(), ttl)
|
||
}
|
||
if stateStore != nil && storeDisabled && sessionHash != "" {
|
||
stateStore.BindSessionConn(groupID, sessionHash, lease.ConnID(), s.openAIWSSessionStickyTTL())
|
||
}
|
||
firstTokenMsValue := -1
|
||
if firstTokenMs != nil {
|
||
firstTokenMsValue = *firstTokenMs
|
||
}
|
||
logOpenAIWSModeDebug(
|
||
"completed account_id=%d conn_id=%s response_id=%s stream=%v duration_ms=%d events=%d token_events=%d terminal_events=%d buffered_events=%d buffered_flushed=%d first_event=%s last_event=%s first_token_ms=%d wrote_downstream=%v client_disconnected=%v",
|
||
account.ID,
|
||
connID,
|
||
truncateOpenAIWSLogValue(strings.TrimSpace(responseID), openAIWSIDValueMaxLen),
|
||
reqStream,
|
||
time.Since(startTime).Milliseconds(),
|
||
eventCount,
|
||
tokenEventCount,
|
||
terminalEventCount,
|
||
bufferedEventCount,
|
||
flushedBufferedEventCount,
|
||
truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen),
|
||
truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen),
|
||
firstTokenMsValue,
|
||
wroteDownstream,
|
||
clientDisconnected,
|
||
)
|
||
|
||
result := resultWithUsage()
|
||
result.ImageCount = imageCounter.Count()
|
||
result.ImageOutputSizes = imageCounter.Sizes()
|
||
return result, nil
|
||
}
|
||
|
||
// ProxyResponsesWebSocketFromClient 处理客户端入站 WebSocket（OpenAI Responses WS Mode）并转发到上游。
|
||
// 当前实现按“单请求 -> 终止事件 -> 下一请求”的顺序代理，适配 Codex CLI 的 turn 模式。
|
||
// stripCodexSparkImageGenerationToolFromRawPayload removes the image_generation
|
||
// tool from a raw /responses payload when the upstream model is gpt-5.3-codex-spark.
|
||
// Spark rejects that tool upstream with HTTP 400 (invalid_request_error, param=tools);
|
||
// Codex clients advertise it by default. Returns the (possibly unchanged) payload,
|
||
// whether it changed, and any JSON decode error.
|
||
func stripCodexSparkImageGenerationToolFromRawPayload(payload []byte, model string) ([]byte, bool, error) {
|
||
if !isCodexSparkModel(model) {
|
||
return payload, false, nil
|
||
}
|
||
return stripOpenAIImageGenerationToolsFromRawPayload(payload)
|
||
}
|