Add WebSocket relay support for /v1/responses

controller/relay.go

@@ -1,11 +1,13 @@
 package controller
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
 	"log"
 	"net/http"
+	"strconv"
 	"strings"
 	"time"
 
@@ -22,6 +24,7 @@ import (
 	"github.com/QuantumNous/new-api/service"
 	"github.com/QuantumNous/new-api/setting"
 	"github.com/QuantumNous/new-api/setting/operation_setting"
+	"github.com/QuantumNous/new-api/setting/ratio_setting"
 	"github.com/QuantumNous/new-api/types"
 
 	"github.com/bytedance/gopkg/util/gopool"
@@ -69,20 +72,38 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 	requestId := c.GetString(common.RequestIdKey)
 	//group := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
 	//originalModel := common.GetContextKeyString(c, constant.ContextKeyOriginalModel)
+	responsesWs := relayFormat == types.RelayFormatOpenAIResponses && websocket.IsWebSocketUpgrade(c.Request)
 
 	var (
 		newAPIError *types.NewAPIError
+		request     dto.Request
 		ws          *websocket.Conn
 	)
 
-	if relayFormat == types.RelayFormatOpenAIRealtime {
+	if relayFormat == types.RelayFormatOpenAIRealtime || responsesWs {
 		var err error
 		ws, err = upgrader.Upgrade(c.Writer, c.Request, nil)
 		if err != nil {
 			helper.WssError(c, ws, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError())
 			return
 		}
 		defer ws.Close()
+
+		if responsesWs {
+			request, err = readFirstResponsesWSRequest(ws)
+			if err != nil {
+				helper.WssError(c, ws, types.NewError(err, types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()).ToOpenAIError())
+				return
+			}
+			if responsesReq, ok := request.(*dto.OpenAIResponsesRequest); ok {
+				c.Set("first_wss_request", responsesReq)
+				newAPIError = setupResponsesWSChannel(c, responsesReq.Model)
+				if newAPIError != nil {
+					helper.WssError(c, ws, newAPIError.ToOpenAIError())
+					return
+				}
+			}
+		}
 	}
 
 	defer func() {
@@ -92,6 +113,14 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 			switch relayFormat {
 			case types.RelayFormatOpenAIRealtime:
 				helper.WssError(c, ws, newAPIError.ToOpenAIError())
+			case types.RelayFormatOpenAIResponses:
+				if responsesWs {
+					helper.WssError(c, ws, newAPIError.ToOpenAIError())
+					return
+				}
+				c.JSON(newAPIError.StatusCode, gin.H{
+					"error": newAPIError.ToOpenAIError(),
+				})
 			case types.RelayFormatClaude:
 				c.JSON(newAPIError.StatusCode, gin.H{
 					"type":  "error",
@@ -105,22 +134,28 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 		}
 	}()
 
-	request, err := helper.GetAndValidateRequest(c, relayFormat)
-	if err != nil {
-		// Map "request body too large" to 413 so clients can handle it correctly
-		if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
-			newAPIError = types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
-		} else {
-			newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest)
+	if request == nil {
+		var err error
+		request, err = helper.GetAndValidateRequest(c, relayFormat)
+		if err != nil {
+			// Map "request body too large" to 413 so clients can handle it correctly
+			if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
+				newAPIError = types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
+			} else {
+				newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest)
+			}
+			return
 		}
-		return
 	}
 
 	relayInfo, err := relaycommon.GenRelayInfo(c, relayFormat, request, ws)
 	if err != nil {
 		newAPIError = types.NewError(err, types.ErrorCodeGenRelayInfoFailed)
 		return
 	}
+	if responsesWs {
+		relayInfo.ClientWs = ws
+	}
 
 	needSensitiveCheck := setting.ShouldCheckPromptSensitive()
 	needCountToken := constant.CountToken
@@ -211,6 +246,12 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 		switch relayFormat {
 		case types.RelayFormatOpenAIRealtime:
 			newAPIError = relay.WssHelper(c, relayInfo)
+		case types.RelayFormatOpenAIResponses:
+			if responsesWs {
+				newAPIError = relay.WssResponsesHelper(c, relayInfo)
+				break
+			}
+			newAPIError = relayHandler(c, relayInfo)
 		case types.RelayFormatClaude:
 			newAPIError = relay.ClaudeHelper(c, relayInfo)
 		case types.RelayFormatGemini:
@@ -241,6 +282,92 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 	}
 }
 
+func readFirstResponsesWSRequest(ws *websocket.Conn) (dto.Request, error) {
+	if ws == nil {
+		return nil, errors.New("websocket connection is nil")
+	}
+	_, payload, err := ws.ReadMessage()
+	if err != nil {
+		return nil, fmt.Errorf("read first websocket message failed: %w", err)
+	}
+	request := &dto.OpenAIResponsesRequest{}
+	if err := json.Unmarshal(payload, request); err != nil {
+		return nil, fmt.Errorf("parse first websocket message failed: %w", err)
+	}
+	if request.Model == "" {
+		return nil, errors.New("model is required")
+	}
+	if request.Input == nil {
+		return nil, errors.New("input is required")
+	}
+	return request, nil
+}
+
+func setupResponsesWSChannel(c *gin.Context, modelName string) *types.NewAPIError {
+	if modelName == "" {
+		return types.NewError(errors.New("model is required"), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
+	}
+	if common.GetContextKeyInt(c, constant.ContextKeyChannelId) != 0 &&
+		common.GetContextKeyString(c, constant.ContextKeyOriginalModel) != "" {
+		return nil
+	}
+
+	modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled)
+	if modelLimitEnable {
+		s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit)
+		if !ok {
+			return types.NewErrorWithStatusCode(errors.New("token has no model access"), types.ErrorCodeInvalidRequest, http.StatusForbidden, types.ErrOptionWithSkipRetry())
+		}
+		tokenModelLimit, ok := s.(map[string]bool)
+		if !ok {
+			tokenModelLimit = map[string]bool{}
+		}
+		matchName := ratio_setting.FormatMatchingModelName(modelName)
+		if _, ok := tokenModelLimit[matchName]; !ok {
+			return types.NewErrorWithStatusCode(fmt.Errorf("model %s is forbidden for this token", modelName), types.ErrorCodeInvalidRequest, http.StatusForbidden, types.ErrOptionWithSkipRetry())
+		}
+	}
+
+	var selectedChannel *model.Channel
+	if channelId, ok := common.GetContextKey(c, constant.ContextKeyTokenSpecificChannelId); ok {
+		id, err := strconv.Atoi(channelId.(string))
+		if err != nil {
+			return types.NewErrorWithStatusCode(err, types.ErrorCodeGetChannelFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
+		}
+		channel, err := model.GetChannelById(id, true)
+		if err != nil {
+			return types.NewErrorWithStatusCode(err, types.ErrorCodeGetChannelFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
+		}
+		if channel.Status != common.ChannelStatusEnabled {
+			return types.NewErrorWithStatusCode(errors.New("channel is disabled"), types.ErrorCodeGetChannelFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry())
+		}
+		selectedChannel = channel
+	} else {
+		usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
+		if usingGroup == "" {
+			usingGroup = common.GetContextKeyString(c, constant.ContextKeyUserGroup)
+		}
+		channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(&service.RetryParam{
+			Ctx:        c,
+			ModelName:  modelName,
+			TokenGroup: usingGroup,
+			Retry:      common.GetPointer(0),
+		})
+		if err != nil {
+			return types.NewErrorWithStatusCode(err, types.ErrorCodeGetChannelFailed, http.StatusServiceUnavailable, types.ErrOptionWithSkipRetry())
+		}
+		if channel == nil {
+			return types.NewErrorWithStatusCode(fmt.Errorf("no available channel for model %s", modelName), types.ErrorCodeGetChannelFailed, http.StatusServiceUnavailable, types.ErrOptionWithSkipRetry())
+		}
+		if usingGroup == "auto" && selectGroup != "" {
+			common.SetContextKey(c, constant.ContextKeyAutoGroup, selectGroup)
+		}
+		selectedChannel = channel
+	}
+
+	return middleware.SetupContextForSelectedChannel(c, selectedChannel, modelName)
+}
+
 var upgrader = websocket.Upgrader{
 	Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol
 	CheckOrigin: func(r *http.Request) bool {

dto/openai_request.go

@@ -812,6 +812,7 @@ type WebSearchOptions struct {
 
 // https://platform.openai.com/docs/api-reference/responses/create
 type OpenAIResponsesRequest struct {
+	Type    string          `json:"type,omitempty"`
 	Model   string          `json:"model"`
 	Input   json.RawMessage `json:"input,omitempty"`
 	Include json.RawMessage `json:"include,omitempty"`

middleware/distributor.go

@@ -20,6 +20,7 @@ import (
 	"github.com/QuantumNous/new-api/types"
 
 	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
 )
 
 type ModelRequest struct {
@@ -177,6 +178,15 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
 	var modelRequest ModelRequest
 	shouldSelectChannel := true
 	var err error
+	if c.Request.Method == http.MethodGet &&
+		websocket.IsWebSocketUpgrade(c.Request) &&
+		strings.HasPrefix(c.Request.URL.Path, "/v1/responses") {
+		modelRequest.Model = c.Query("model")
+		if strings.HasPrefix(c.Request.URL.Path, "/v1/responses/compact") && modelRequest.Model != "" {
+			modelRequest.Model = ratio_setting.WithCompactModelSuffix(modelRequest.Model)
+		}
+		return &modelRequest, false, nil
+	}
 	if strings.Contains(c.Request.URL.Path, "/mj/") {
 		relayMode := relayconstant.Path2RelayModeMidjourney(c.Request.URL.Path)
 		if relayMode == relayconstant.RelayModeMidjourneyTaskFetch ||

relay/channel/api_request.go

@@ -356,6 +356,11 @@ func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody
 	if err != nil {
 		return nil, fmt.Errorf("get request url failed: %w", err)
 	}
+	if strings.HasPrefix(fullRequestURL, "https://") {
+		fullRequestURL = "wss://" + strings.TrimPrefix(fullRequestURL, "https://")
+	} else if strings.HasPrefix(fullRequestURL, "http://") {
+		fullRequestURL = "ws://" + strings.TrimPrefix(fullRequestURL, "http://")
+	}
 	targetHeader := http.Header{}
 	err = a.SetupRequestHeader(c, &targetHeader, info)
 	if err != nil {

relay/websocket.go

@@ -1,13 +1,19 @@
 package relay
 
 import (
+	"errors"
 	"fmt"
+	"time"
 
 	"github.com/QuantumNous/new-api/dto"
+	"github.com/QuantumNous/new-api/logger"
+	"github.com/QuantumNous/new-api/relay/channel"
 	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/QuantumNous/new-api/relay/helper"
 	"github.com/QuantumNous/new-api/service"
 	"github.com/QuantumNous/new-api/types"
 
+	"github.com/bytedance/gopkg/util/gopool"
 	"github.com/gin-gonic/gin"
 	"github.com/gorilla/websocket"
 )
@@ -44,3 +50,103 @@ func WssHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.
 	service.PostWssConsumeQuota(c, info, info.UpstreamModelName, usage.(*dto.RealtimeUsage), "")
 	return nil
 }
+
+func WssResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
+	info.InitChannelMeta(c)
+	if info.ClientWs == nil {
+		return types.NewError(errors.New("client websocket connection is nil"), types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
+	}
+
+	adaptor := GetAdaptor(info.ApiType)
+	if adaptor == nil {
+		return types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry())
+	}
+	adaptor.Init(info)
+
+	targetWs, err := channel.DoWssRequest(adaptor, c, info, nil)
+	if err != nil {
+		return types.NewError(err, types.ErrorCodeDoRequestFailed)
+	}
+	info.TargetWs = targetWs
+	defer info.TargetWs.Close()
+
+	if err := sendInitialResponsesWSRequest(c, adaptor, info); err != nil {
+		return types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithSkipRetry())
+	}
+	if err := proxyResponsesWS(c, info.ClientWs, info.TargetWs); err != nil {
+		return types.NewError(err, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
+	}
+	if err := service.SettleBilling(c, info, info.FinalPreConsumedQuota); err != nil {
+		logger.LogError(c, "responses websocket settle billing failed: "+err.Error())
+	}
+	return nil
+}
+
+func sendInitialResponsesWSRequest(c *gin.Context, adaptor channel.Adaptor, info *relaycommon.RelayInfo) error {
+	request, ok := info.Request.(*dto.OpenAIResponsesRequest)
+	if !ok || request == nil {
+		return errors.New("invalid responses websocket request")
+	}
+	initialRequest := *request
+	if initialRequest.Type == "" {
+		initialRequest.Type = "response.create"
+	}
+	converted, err := adaptor.ConvertOpenAIResponsesRequest(c, info, initialRequest)
+	if err != nil {
+		return err
+	}
+	switch v := converted.(type) {
+	case string:
+		return helper.WssString(c, info.TargetWs, v)
+	case []byte:
+		return info.TargetWs.WriteMessage(websocket.TextMessage, v)
+	default:
+		return helper.WssObject(c, info.TargetWs, v)
+	}
+}
+
+func proxyResponsesWS(c *gin.Context, clientConn, targetConn *websocket.Conn) error {
+	errChan := make(chan error, 2)
+	forward := func(src, dst *websocket.Conn, direction string) {
+		defer func() {
+			if r := recover(); r != nil {
+				errChan <- fmt.Errorf("%s panic: %v", direction, r)
+			}
+		}()
+		for {
+			messageType, payload, err := src.ReadMessage()
+			if err != nil {
+				if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
+					errChan <- nil
+					return
+				}
+				errChan <- fmt.Errorf("%s read failed: %w", direction, err)
+				return
+			}
+			if err := dst.WriteMessage(messageType, payload); err != nil {
+				errChan <- fmt.Errorf("%s write failed: %w", direction, err)
+				return
+			}
+		}
+	}
+
+	gopool.Go(func() {
+		forward(clientConn, targetConn, "client->target")
+	})
+	gopool.Go(func() {
+		forward(targetConn, clientConn, "target->client")
+	})
+
+	select {
+	case <-c.Request.Context().Done():
+		return nil
+	case err := <-errChan:
+		if err == nil {
+			deadline := time.Now().Add(500 * time.Millisecond)
+			_ = clientConn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), deadline)
+			_ = targetConn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), deadline)
+			return nil
+		}
+		return err
+	}
+}

router/relay-router.go

@@ -78,6 +78,9 @@ func SetRelayRouter(router *gin.Engine) {
 		wsRouter.GET("/realtime", func(c *gin.Context) {
 			controller.Relay(c, types.RelayFormatOpenAIRealtime)
 		})
+		wsRouter.GET("/responses", func(c *gin.Context) {
+			controller.Relay(c, types.RelayFormatOpenAIResponses)
+		})
 	}
 	{
 		//http router