cneicy/new-api-custom
1
0 Fork 0
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

feat: publish sanitized custom.97 source

Browse Source
This commit is contained in:
cneicy committed 2026-10-27 20:43:39 +08:00
1 parent 8e6ed0317b
commit 86c0276156
1022 files changed
+69903 -19207
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

relay/responses_websocket.go
+85 -71
View File
@@ -79,11 +79,13 @@ type responsesWSControl struct {
// Only the request worker reads or changes billing state. Socket readers pass
// bounded messages to it; cancellation never performs an independent refund.
type responsesWSCallState struct {
inbox chan responsesWSMessage
controls chan responsesWSControl
done chan struct{}
terminal *responsesWSMessage
closeAfter bool
inbox chan responsesWSMessage
controls chan responsesWSControl
done chan struct{}
terminal *responsesWSMessage
// A nonzero closeCode ends the connection after this request.
closeCode int
closeReason string
}
type responsesWSSession struct {
@@ -95,6 +97,7 @@ type responsesWSSession struct {
requestID string
nextEventIndex int
workers sync.WaitGroup
closeOnce sync.Once
clientWriteMu sync.Mutex
targetWriteMu sync.Mutex
@@ -129,7 +132,9 @@ func ResponsesWebSocketHelper(c *gin.Context, client *websocket.Conn, runner Res
}
client.SetReadLimit(int64(maxMB) << 20)
defer func() {
s.shutdown()
// Unless an earlier shutdown ended the read, the client closed or left.
// Gorilla already answered a client close frame, so this tells the upstream.
s.shutdown(websocket.CloseNormalClosure, "")
s.workers.Wait()
}()
@@ -176,17 +181,20 @@ func ResponsesWebSocketHelper(c *gin.Context, client *websocket.Conn, runner Res
func (s *responsesWSSession) runRequest(state *responsesWSCallState, message []byte, envelope responsesWSCreateEvent, streamID string, requestID string) {
create, parseErr := normalizeResponsesWSCreateEvent(message, envelope, streamID)
var apiErr *types.NewAPIError
var upstreamURLs []string
defer func() {
if recovered := recover(); recovered != nil {
apiErr = types.NewError(fmt.Errorf("responses websocket request panic: %v", recovered), types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
state.closeAfter = true
state.closeCode = websocket.CloseInternalServerErr
}
// Finish the middleware count before publishing the terminal event. Hold
// admission while writing it so an immediate next create cannot race
// the current request's release; socket close needs neither lock.
// the current request's release. A failing request also shuts down
// before that release, so a concurrent upstream close cannot report it
// as a normal closure; shutdown itself takes neither lock.
outgoing := state.terminal
if apiErr != nil {
if body, err := buildResponsesWSErrorPayload(envelope.EventID, streamID, apiErr); err == nil {
if body, err := buildResponsesWSErrorPayload(envelope.EventID, streamID, apiErr, upstreamURLs...); err == nil {
outgoing = &responsesWSMessage{kind: websocket.TextMessage, body: body}
}
}
@@ -194,18 +202,18 @@ func (s *responsesWSSession) runRequest(state *responsesWSCallState, message []b
s.stateMu.Lock()
if outgoing != nil {
if err := s.client.SetWriteDeadline(time.Now().Add(responsesWSWriteTimeout)); err != nil {
state.closeAfter = true
state.closeCode = websocket.CloseInternalServerErr
} else if err := s.client.WriteMessage(outgoing.kind, outgoing.body); err != nil {
state.closeAfter = true
state.closeCode = websocket.CloseInternalServerErr
}
}
if state.closeCode != 0 {
s.shutdown(state.closeCode, state.closeReason)
}
s.current = nil
close(state.done)
s.stateMu.Unlock()
s.clientWriteMu.Unlock()
if state.closeAfter {
s.shutdown()
}
}()
request := s.request.Clone(s.ctx)
request.Body = io.NopCloser(bytes.NewReader(message))
@@ -217,10 +225,12 @@ func (s *responsesWSSession) runRequest(state *responsesWSCallState, message []b
}
c.Request.Body = io.NopCloser(bytes.NewReader(create.Body))
c.Request.ContentLength = int64(len(create.Body))
return s.runCall(c, state, create)
callErr := s.runCall(c, state, create)
upstreamURLs = common.UpstreamAddresses(c)
return callErr
})
if apiErr != nil && (apiErr.StatusCode == http.StatusUnauthorized || apiErr.StatusCode == http.StatusForbidden) {
state.closeAfter = true
state.closeCode = websocket.ClosePolicyViolation
}
}
@@ -233,7 +243,7 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
defer func() {
if recovered := recover(); recovered != nil {
apiErr = types.NewError(fmt.Errorf("responses websocket call panic: %v", recovered), types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
state.closeAfter = true
state.closeCode = websocket.CloseInternalServerErr
}
if info == nil && modelName != "" {
info = &relaycommon.RelayInfo{OriginModelName: modelName, UsingGroup: common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup), StartTime: started}
@@ -276,7 +286,7 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
return apiErr
}
if err := s.writeTarget(websocket.TextMessage, payload); err != nil {
state.closeAfter = true
state.closeCode = websocket.CloseInternalServerErr
return types.NewError(err, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
}
} else {
@@ -337,7 +347,7 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
return types.NewError(context.Canceled, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
}
if err := s.writeTarget(websocket.TextMessage, payload); err != nil {
state.closeAfter = true
state.closeCode = websocket.CloseInternalServerErr
return types.NewError(err, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
}
s.lockedModel, s.lockedChannelID, s.lockedGroup = modelName, channel.Id, info.UsingGroup
@@ -350,6 +360,7 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
appconstant.ContextKeyChannelParamOverride, appconstant.ContextKeyChannelHeaderOverride, appconstant.ContextKeyChannelOrganization,
appconstant.ContextKeyChannelAutoBan, appconstant.ContextKeyChannelModelMapping, appconstant.ContextKeyChannelStatusCodeMapping,
appconstant.ContextKeyChannelIsMultiKey, appconstant.ContextKeyChannelMultiKeyIndex, appconstant.ContextKeyChannelKey, appconstant.ContextKeyChannelBaseUrl,
appconstant.ContextKeyUpstreamRequestOrigin,
} {
if value, ok := c.Get(string(key)); ok {
s.lockedContext[key] = value
@@ -385,7 +396,7 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
idle.Reset(timeout)
if incoming.err != nil {
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonScannerErr, incoming.err)
state.closeAfter = true
state.closeCode, state.closeReason = websocket.CloseInternalServerErr, "upstream connection closed"
ConsumeResponsesQuota(c, info, accumulator.Finish())
return nil
}
@@ -399,7 +410,7 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
} else {
if event.Type != "error" && event.StreamID != "" && event.StreamID != create.StreamID {
if err := s.writeClient(incoming.kind, incoming.body); err != nil {
s.shutdown()
s.shutdown(websocket.CloseInternalServerErr, "")
}
continue
}
@@ -417,7 +428,7 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
terminal, ambiguous, controlError := responsesWSErrorEndsRequest(rejection, create.StreamID, responseID, sentControl)
if !terminal {
if err := s.writeClient(incoming.kind, incoming.body); err != nil {
s.shutdown()
s.shutdown(websocket.CloseInternalServerErr, "")
}
// Only a control error in this stream resolves its pending control.
if controlError {
@@ -433,10 +444,13 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
}
info.StreamStatus.MarkFailed(code, rejection.Error.Type, rejection.Status)
}
accumulator.Observe(&event.ResponsesStreamResponse)
accumulator.Observe(&event.ResponsesStreamResponse, incoming.body)
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonDone, nil)
s.lastResponseID = responseID
state.terminal, state.closeAfter = &incoming, ambiguous
state.terminal = &incoming
if ambiguous {
state.closeCode = websocket.CloseInternalServerErr
}
ConsumeResponsesQuota(c, info, accumulator.Finish())
return nil
}
@@ -461,7 +475,7 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
responseID = event.Response.ID
}
}
accumulator.Observe(&event.ResponsesStreamResponse)
accumulator.Observe(&event.ResponsesStreamResponse, incoming.body)
}
switch event.Type {
case "response.completed", "response.done", "response.incomplete", "response.failed", "response.cancelled", "response.canceled":
@@ -474,11 +488,11 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
return nil
}
if err := s.writeClient(incoming.kind, incoming.body); err != nil {
s.shutdown()
s.shutdown(websocket.CloseInternalServerErr, "")
}
if accepted && pendingControl != nil {
if err := s.writeTarget(websocket.TextMessage, pendingControl); err != nil {
s.shutdown()
s.shutdown(websocket.CloseInternalServerErr, "")
}
sentControl = pendingControl
pendingControl = nil
@@ -493,12 +507,12 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
continue
}
if err := s.writeTarget(websocket.TextMessage, control.body); err != nil {
s.shutdown()
s.shutdown(websocket.CloseInternalServerErr, "")
}
sentControl = control.body
case <-idle.C:
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, context.DeadlineExceeded)
state.closeAfter = true
state.closeCode, state.closeReason = websocket.CloseInternalServerErr, "upstream response timed out"
ConsumeResponsesQuota(c, info, accumulator.Finish())
return nil
case <-s.ctx.Done():
@@ -642,27 +656,27 @@ func (s *responsesWSSession) startTargetReader(target *websocket.Conn) {
if state := s.getCurrent(); state != nil {
select {
case state.inbox <- incoming:
case <-state.done:
if err == nil {
if writeErr := s.writeClient(kind, body); writeErr != nil {
s.shutdown()
return
}
} else {
s.shutdown()
if err != nil {
return
}
continue
case <-state.done:
case <-s.ctx.Done():
return
}
} else if err == nil {
if writeErr := s.writeClient(kind, body); writeErr != nil {
s.shutdown()
return
}
} else {
s.shutdown()
}
// No request owns this message. An idle upstream that closed
// normally ends the client session with the same code.
if err != nil {
code := websocket.CloseInternalServerErr
if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
code = err.(*websocket.CloseError).Code
}
s.shutdown(code, "upstream connection closed")
return
}
if writeErr := s.writeClient(kind, body); writeErr != nil {
s.shutdown(websocket.CloseInternalServerErr, "")
return
}
}
@@ -695,6 +709,7 @@ func (s *responsesWSSession) setTarget(target *websocket.Conn) bool {
s.connectionMu.Lock()
defer s.connectionMu.Unlock()
if s.ctx.Err() != nil {
_ = target.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))
_ = target.Close()
return false
}
@@ -734,23 +749,28 @@ func (s *responsesWSSession) sendError(eventID, streamID string, apiErr *types.N
}
}
func (s *responsesWSSession) closeTarget() {
s.connectionMu.Lock()
target, unregister := s.target, s.unregister
s.target, s.unregister = nil, nil
s.connectionMu.Unlock()
if unregister != nil {
unregister()
}
if target != nil {
_ = target.Close()
}
}
func (s *responsesWSSession) shutdown() {
s.cancel()
s.closeTarget()
_ = s.client.Close()
// shutdown ends the connection once. Each peer gets a close frame that says
// why, bounded to one second, and both sockets close even if a write failed.
// Close frames are control writes, which gorilla allows alongside a blocked
// data writer, so shutdown takes no write lock and never waits for workers.
func (s *responsesWSSession) shutdown(code int, reason string) {
s.closeOnce.Do(func() {
s.cancel()
s.connectionMu.Lock()
target, unregister := s.target, s.unregister
s.target, s.unregister = nil, nil
s.connectionMu.Unlock()
if unregister != nil {
unregister()
}
closeMessage := websocket.FormatCloseMessage(code, reason)
_ = s.client.WriteControl(websocket.CloseMessage, closeMessage, time.Now().Add(time.Second))
if target != nil {
_ = target.WriteControl(websocket.CloseMessage, closeMessage, time.Now().Add(time.Second))
_ = target.Close()
}
_ = s.client.Close()
})
}
func (s *responsesWSSession) registerChannelClose(channelID int) {
@@ -766,15 +786,9 @@ func (s *responsesWSSession) registerChannelClose(channelID int) {
}
// closeForPolicy tells both peers why the connection ends, like the realtime
// relay does; WriteControl is safe alongside a blocked data writer.
// relay does.
func (s *responsesWSSession) closeForPolicy(reason string) {
closeMessage := websocket.FormatCloseMessage(websocket.ClosePolicyViolation, reason)
deadline := time.Now().Add(time.Second)
_ = s.client.WriteControl(websocket.CloseMessage, closeMessage, deadline)
if target := s.getTarget(); target != nil {
_ = target.WriteControl(websocket.CloseMessage, closeMessage, deadline)
}
s.shutdown()
s.shutdown(websocket.ClosePolicyViolation, reason)
}
// Stream identity belongs to the WebSocket envelope. For the legacy wrapped
@@ -886,7 +900,7 @@ func buildResponsesWSCreateEvent(jsonData []byte, generate common.RawMessage, st
return common.Marshal(event)
}
func buildResponsesWSErrorPayload(eventID, streamID string, apiErr *types.NewAPIError) ([]byte, error) {
func buildResponsesWSErrorPayload(eventID, streamID string, apiErr *types.NewAPIError, upstreamURLs ...string) ([]byte, error) {
if apiErr == nil {
return nil, errors.New("api error is nil")
}
@@ -894,7 +908,7 @@ func buildResponsesWSErrorPayload(eventID, streamID string, apiErr *types.NewAPI
if status == 0 {
status = http.StatusInternalServerError
}
openaiErr := apiErr.ToOpenAIError()
openaiErr := apiErr.ToOpenAIError(upstreamURLs...)
return common.Marshal(&responsesWSErrorEvent{
Type: "error",
Status: status,