388 lines
11 KiB
Go
388 lines
11 KiB
Go
package service
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
const (
|
|
openAIWSClientReadLimitBytesDefault int64 = 64 * 1024 * 1024
|
|
openAIWSHTTPBridgeThresholdBytesDefault int64 = 15 * 1024 * 1024
|
|
openAIWSHTTPBridgeErrorBodyLimitBytes = 64 * 1024
|
|
)
|
|
|
|
func ResolveOpenAIWSClientReadLimitBytes(cfg *config.Config) int64 {
|
|
if cfg == nil || cfg.Gateway.OpenAIWS.ClientReadLimitBytes <= 0 {
|
|
return openAIWSClientReadLimitBytesDefault
|
|
}
|
|
return cfg.Gateway.OpenAIWS.ClientReadLimitBytes
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) openAIWSHTTPBridgeEnabled() bool {
|
|
return s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.HTTPBridgeEnabled
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) openAIWSHTTPBridgeThresholdBytes() int64 {
|
|
if s == nil || s.cfg == nil || s.cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes <= 0 {
|
|
return openAIWSHTTPBridgeThresholdBytesDefault
|
|
}
|
|
return s.cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) shouldBridgeOpenAIWSHTTP(payloadBytes int, previousResponseID string) bool {
|
|
if !s.openAIWSHTTPBridgeEnabled() {
|
|
return false
|
|
}
|
|
if strings.TrimSpace(previousResponseID) != "" {
|
|
return false
|
|
}
|
|
threshold := s.openAIWSHTTPBridgeThresholdBytes()
|
|
return threshold > 0 && int64(payloadBytes) >= threshold
|
|
}
|
|
|
|
func prepareOpenAIWSHTTPBridgeBody(payload []byte) ([]byte, error) {
|
|
var body map[string]any
|
|
if err := json.Unmarshal(payload, &body); err != nil {
|
|
return nil, err
|
|
}
|
|
if body == nil {
|
|
return nil, errors.New("response.create payload must be a JSON object")
|
|
}
|
|
delete(body, "type")
|
|
delete(body, "generate")
|
|
delete(body, "previous_response_id")
|
|
body["stream"] = true
|
|
return json.Marshal(body)
|
|
}
|
|
|
|
type openAIWSToolCallReplayCollector struct {
|
|
items []json.RawMessage
|
|
seen map[string]struct{}
|
|
}
|
|
|
|
func (c *openAIWSToolCallReplayCollector) AddEvent(eventType string, message []byte) {
|
|
switch strings.TrimSpace(eventType) {
|
|
case "response.output_item.done":
|
|
c.addItem(gjson.GetBytes(message, "item"))
|
|
case "response.completed", "response.done":
|
|
output := gjson.GetBytes(message, "response.output")
|
|
if !output.IsArray() {
|
|
return
|
|
}
|
|
for _, item := range output.Array() {
|
|
c.addItem(item)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *openAIWSToolCallReplayCollector) Items() []json.RawMessage {
|
|
return cloneOpenAIWSRawMessages(c.items)
|
|
}
|
|
|
|
func (c *openAIWSToolCallReplayCollector) addItem(item gjson.Result) {
|
|
if !item.Exists() || item.Type != gjson.JSON {
|
|
return
|
|
}
|
|
raw := strings.TrimSpace(item.Raw)
|
|
if raw == "" || !strings.HasPrefix(raw, "{") {
|
|
return
|
|
}
|
|
if !isCodexToolCallContextItemType(item.Get("type").String()) {
|
|
return
|
|
}
|
|
key := strings.TrimSpace(item.Get("id").String())
|
|
if key == "" {
|
|
key = strings.TrimSpace(item.Get("call_id").String())
|
|
}
|
|
if key == "" {
|
|
key = raw
|
|
}
|
|
if c.seen == nil {
|
|
c.seen = make(map[string]struct{})
|
|
}
|
|
if _, ok := c.seen[key]; ok {
|
|
return
|
|
}
|
|
c.seen[key] = struct{}{}
|
|
c.items = append(c.items, json.RawMessage(raw))
|
|
}
|
|
|
|
func buildOpenAIWSHTTPBridgeErrorEvent(statusCode int, message string) []byte {
|
|
message = strings.TrimSpace(message)
|
|
if message == "" {
|
|
message = http.StatusText(statusCode)
|
|
}
|
|
if message == "" {
|
|
message = "upstream request failed"
|
|
}
|
|
event := map[string]any{
|
|
"type": "error",
|
|
"status": statusCode,
|
|
"error": map[string]any{
|
|
"type": "upstream_error",
|
|
"message": message,
|
|
},
|
|
}
|
|
body, err := json.Marshal(event)
|
|
if err != nil {
|
|
return []byte(`{"type":"error","error":{"type":"upstream_error","message":"upstream request failed"}}`)
|
|
}
|
|
return body
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
|
ctx context.Context,
|
|
c *gin.Context,
|
|
account *Account,
|
|
token string,
|
|
payload []byte,
|
|
payloadBytes int,
|
|
originalModel string,
|
|
imageBillingModel string,
|
|
imageSizeTier string,
|
|
imageInputSize string,
|
|
turn int,
|
|
writeClientMessage func([]byte) error,
|
|
) (*OpenAIForwardResult, error) {
|
|
if s == nil {
|
|
return nil, errors.New("service is nil")
|
|
}
|
|
if s.httpUpstream == nil {
|
|
return nil, errors.New("openai http upstream is nil")
|
|
}
|
|
if account == nil {
|
|
return nil, errors.New("account is nil")
|
|
}
|
|
if writeClientMessage == nil {
|
|
return nil, errors.New("client websocket writer is nil")
|
|
}
|
|
|
|
body, err := prepareOpenAIWSHTTPBridgeBody(payload)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("prepare http bridge body: %w", err)
|
|
}
|
|
|
|
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
|
|
upstreamReq, err := s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token)
|
|
releaseUpstreamCtx()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
proxyURL := ""
|
|
if account.ProxyID != nil && account.Proxy != nil {
|
|
proxyURL = account.Proxy.URL()
|
|
}
|
|
if c != nil {
|
|
c.Set("openai_passthrough", true)
|
|
c.Set("openai_ws_http_bridge", true)
|
|
}
|
|
|
|
turnStart := time.Now()
|
|
resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
|
|
if err != nil {
|
|
safeErr := sanitizeUpstreamErrorMessage(err.Error())
|
|
_ = writeClientMessage(buildOpenAIWSHTTPBridgeErrorEvent(http.StatusBadGateway, "Upstream request failed"))
|
|
return nil, fmt.Errorf("upstream http bridge request failed: %s", safeErr)
|
|
}
|
|
defer func() { _ = resp.Body.Close() }()
|
|
|
|
if resp.StatusCode >= 400 {
|
|
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, openAIWSHTTPBridgeErrorBodyLimitBytes))
|
|
upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
|
|
if upstreamMsg == "" {
|
|
upstreamMsg = http.StatusText(resp.StatusCode)
|
|
}
|
|
_ = writeClientMessage(buildOpenAIWSHTTPBridgeErrorEvent(resp.StatusCode, upstreamMsg))
|
|
return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg)
|
|
}
|
|
|
|
responseID := ""
|
|
usage := OpenAIUsage{}
|
|
imageCounter := newOpenAIImageOutputCounter()
|
|
var firstTokenMs *int
|
|
reqStream := openAIWSPayloadBoolFromRaw(body, "stream", true)
|
|
eventCount := 0
|
|
tokenEventCount := 0
|
|
terminalEventCount := 0
|
|
replayCollector := &openAIWSToolCallReplayCollector{}
|
|
firstEventType := ""
|
|
lastEventType := ""
|
|
sawDone := false
|
|
wroteDownstream := false
|
|
clientDisconnected := false
|
|
mappedModel := ""
|
|
needModelReplace := false
|
|
var mappedModelBytes []byte
|
|
if originalModel != "" {
|
|
mappedModel = normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel))
|
|
needModelReplace = mappedModel != "" && mappedModel != originalModel
|
|
if needModelReplace {
|
|
mappedModelBytes = []byte(mappedModel)
|
|
}
|
|
}
|
|
|
|
resultWithUsage := func() *OpenAIForwardResult {
|
|
imageCount := imageCounter.Count()
|
|
result := &OpenAIForwardResult{
|
|
RequestID: responseID,
|
|
Usage: usage,
|
|
Model: originalModel,
|
|
UpstreamModel: mappedModel,
|
|
ServiceTier: extractOpenAIServiceTierFromBody(body),
|
|
ReasoningEffort: extractOpenAIReasoningEffortFromBody(body, originalModel),
|
|
Stream: reqStream,
|
|
OpenAIWSMode: true,
|
|
ResponseHeaders: cloneHeader(resp.Header),
|
|
Duration: time.Since(turnStart),
|
|
FirstTokenMs: firstTokenMs,
|
|
}
|
|
if replayInput := replayCollector.Items(); len(replayInput) > 0 {
|
|
result.wsReplayInput = replayInput
|
|
result.wsReplayInputExists = true
|
|
}
|
|
if imageCount > 0 {
|
|
result.ImageCount = imageCount
|
|
result.ImageSize = imageSizeTier
|
|
result.ImageInputSize = imageInputSize
|
|
result.ImageOutputSizes = imageCounter.Sizes()
|
|
result.BillingModel = imageBillingModel
|
|
}
|
|
return result
|
|
}
|
|
|
|
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)
|
|
defer putSSEScannerBuf64K(scanBuf)
|
|
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
data, ok := extractOpenAISSEDataLine(line)
|
|
if !ok {
|
|
continue
|
|
}
|
|
trimmedData := strings.TrimSpace(data)
|
|
if trimmedData == "" {
|
|
continue
|
|
}
|
|
if trimmedData == "[DONE]" {
|
|
sawDone = true
|
|
continue
|
|
}
|
|
|
|
upstreamMessage := []byte(trimmedData)
|
|
eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(upstreamMessage)
|
|
if responseID == "" && eventResponseID != "" {
|
|
responseID = eventResponseID
|
|
}
|
|
if eventType != "" {
|
|
eventCount++
|
|
if firstEventType == "" {
|
|
firstEventType = eventType
|
|
}
|
|
lastEventType = eventType
|
|
}
|
|
if isOpenAIWSTokenEvent(eventType) {
|
|
tokenEventCount++
|
|
if firstTokenMs == nil {
|
|
ms := int(time.Since(turnStart).Milliseconds())
|
|
firstTokenMs = &ms
|
|
}
|
|
}
|
|
if openAIWSEventShouldParseUsage(eventType) {
|
|
parseOpenAIWSResponseUsageFromCompletedEvent(upstreamMessage, &usage)
|
|
}
|
|
imageCounter.AddSSEData(upstreamMessage)
|
|
|
|
if needModelReplace && len(mappedModelBytes) > 0 && openAIWSEventMayContainModel(eventType) && strings.Contains(trimmedData, mappedModel) {
|
|
upstreamMessage = replaceOpenAIWSMessageModel(upstreamMessage, mappedModel, originalModel)
|
|
}
|
|
if s.toolCorrector != nil && openAIWSEventMayContainToolCalls(eventType) && openAIWSMessageLikelyContainsToolCalls(upstreamMessage) {
|
|
if corrected, changed := s.toolCorrector.CorrectToolCallsInSSEBytes(upstreamMessage); changed {
|
|
upstreamMessage = corrected
|
|
}
|
|
}
|
|
replayCollector.AddEvent(eventType, upstreamMessage)
|
|
|
|
if !clientDisconnected {
|
|
if err := writeClientMessage(upstreamMessage); err != nil {
|
|
if isOpenAIWSClientDisconnectError(err) {
|
|
clientDisconnected = true
|
|
closeStatus, closeReason := summarizeOpenAIWSReadCloseError(err)
|
|
logOpenAIWSModeInfo(
|
|
"ingress_ws_http_bridge_client_disconnected_drain account_id=%d turn=%d close_status=%s close_reason=%s",
|
|
account.ID,
|
|
turn,
|
|
closeStatus,
|
|
truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen),
|
|
)
|
|
} else {
|
|
return nil, wrapOpenAIWSIngressTurnError(
|
|
"write_client",
|
|
fmt.Errorf("write client websocket event: %w", err),
|
|
wroteDownstream,
|
|
)
|
|
}
|
|
} else {
|
|
wroteDownstream = true
|
|
}
|
|
}
|
|
|
|
if eventType == "error" {
|
|
errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage)
|
|
s.persistOpenAIWSRateLimitSignal(ctx, account, resp.Header, upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw)
|
|
errMessage := strings.TrimSpace(errMsgRaw)
|
|
if errMessage == "" {
|
|
errMessage = "upstream error event"
|
|
}
|
|
return resultWithUsage(), errors.New(errMessage)
|
|
}
|
|
if isOpenAIWSTerminalEvent(eventType) {
|
|
terminalEventCount++
|
|
firstTokenMsValue := -1
|
|
if firstTokenMs != nil {
|
|
firstTokenMsValue = *firstTokenMs
|
|
}
|
|
logOpenAIWSModeInfo(
|
|
"ingress_ws_http_bridge_turn_completed account_id=%d turn=%d response_id=%s payload_bytes=%d duration_ms=%d events=%d token_events=%d terminal_events=%d first_event=%s last_event=%s first_token_ms=%d client_disconnected=%v",
|
|
account.ID,
|
|
turn,
|
|
truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen),
|
|
payloadBytes,
|
|
time.Since(turnStart).Milliseconds(),
|
|
eventCount,
|
|
tokenEventCount,
|
|
terminalEventCount,
|
|
truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen),
|
|
truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen),
|
|
firstTokenMsValue,
|
|
clientDisconnected,
|
|
)
|
|
return resultWithUsage(), nil
|
|
}
|
|
}
|
|
if err := scanner.Err(); err != nil {
|
|
return resultWithUsage(), fmt.Errorf("read upstream http bridge stream: %w", err)
|
|
}
|
|
if sawDone && eventCount > 0 {
|
|
return resultWithUsage(), nil
|
|
}
|
|
return resultWithUsage(), errors.New("upstream http bridge stream ended before terminal event")
|
|
}
|