diff --git a/.gitignore b/.gitignore index dc328dd6c80c..5bb21bb2a84e 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,14 @@ upload build *.db-journal logs + +# Local planning/review artifacts generated by agent workflows +.code-forge/ +.planning/ + +# Local architecture specs/plans +code-graph-report.md +docs/superpowers/ web/dist web/node_modules .env diff --git a/model/channel.go b/model/channel.go index 0f8cdb101ec8..f8b742243792 100644 --- a/model/channel.go +++ b/model/channel.go @@ -968,6 +968,9 @@ func (channel *Channel) ValidateSettings() error { if err := channelParams.ValidateHTTPTransport(); err != nil { return err } + if err := channelParams.ValidateForceUpstreamStream(); err != nil { + return err + } channelOtherSettings := &dto.ChannelOtherSettings{} if channel.OtherSettings != "" { err := common.UnmarshalJsonStr(channel.OtherSettings, channelOtherSettings) diff --git a/relay/channel/openai/adaptor.go b/relay/channel/openai/adaptor.go index 64ae3102b3c2..991b9e2edef4 100644 --- a/relay/channel/openai/adaptor.go +++ b/relay/channel/openai/adaptor.go @@ -36,11 +36,15 @@ import ( "github.com/gin-gonic/gin" ) +// Adaptor implements the OpenAI-compatible channel adaptor, handling request +// conversion, header setup, and response dispatch for OpenAI, Azure, and +// other OpenAI-compatible upstreams. type Adaptor struct { ChannelType int ResponseFormat string } +// ConvertGeminiRequest converts a Gemini chat request to the upstream request body. func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error) { result, err := service.ConvertRequest(c, info, types.RelayFormatOpenAI, request) if err != nil { @@ -53,6 +57,7 @@ func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayIn return a.ConvertOpenAIRequest(c, info, openaiRequest) } +// ConvertClaudeRequest converts a Claude request to the upstream request body. func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.ClaudeRequest) (any, error) { //if !strings.Contains(request.Model, "claude") { // return nil, fmt.Errorf("you are using openai channel type with path /v1/messages, only claude model supported convert, but got %s", request.Model) @@ -72,6 +77,12 @@ func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayIn if !ok { return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value) } + // Preserve the original stream flag from the Claude request. The format + // converter may not carry it over, and ConvertOpenAIRequest needs it to + // set info.IsStream correctly for DoResponse routing. + if request.Stream != nil { + aiRequest.Stream = request.Stream + } //if common.DebugEnabled { // println(fmt.Sprintf("convert claude to openai request result: %s", common.GetJsonString(aiRequest))) // // Save request body to file for debugging @@ -89,6 +100,7 @@ func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayIn return a.ConvertOpenAIRequest(c, info, aiRequest) } +// Init initializes the adaptor with channel metadata from RelayInfo. func (a *Adaptor) Init(info *relaycommon.RelayInfo) { a.ChannelType = info.ChannelType @@ -102,6 +114,7 @@ func (a *Adaptor) Init(info *relaycommon.RelayInfo) { } } +// GetRequestURL returns the upstream endpoint URL based on relay mode and channel configuration. func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { if info.RelayMode == relayconstant.RelayModeRealtime { if strings.HasPrefix(info.ChannelBaseUrl, "https://") { @@ -180,6 +193,7 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { } } +// SetupRequestHeader sets authentication and routing headers on the upstream request. func (a *Adaptor) SetupRequestHeader(c *gin.Context, header *http.Header, info *relaycommon.RelayInfo) error { channel.SetupApiRequestHeader(info, c, header) if info.ChannelType == constant.ChannelTypeAzure { @@ -241,11 +255,54 @@ func (a *Adaptor) SetupRequestHeader(c *gin.Context, header *http.Header, info * return nil } +// ConvertOpenAIRequest transforms a client-side GeneralOpenAIRequest into the +// upstream-specific request body. When the channel has ForceUpstreamStream +// enabled and the client requested non-streaming, it forces stream=true on the +// upstream request and sets UpstreamStreamForced so DoResponse routes through +// the buffered SSE aggregation handler. func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) { if request == nil { return nil, errors.New("request is nil") } - if info.ChannelType != constant.ChannelTypeOpenAI && info.ChannelType != constant.ChannelTypeAzure { + // Reset forced-stream flags on each retry attempt. RelayInfo is reused + // across retries (controller/relay.go), so flags set by a previous + // channel must not leak into the current one. + info.IsStream = lo.FromPtrOr(request.Stream, false) + info.UpstreamStreamForced = false + // Force upstream streaming when channel requests it and client asked for non-stream. + // The SSE response will be aggregated by OaiBufferedStreamHandler in DoResponse. + // Do NOT set info.IsStream here -- DoApiRequest uses it to set SSE headers + // and start a ping goroutine for the downstream client, which would corrupt + // the non-streaming JSON response. DoResponse routes on UpstreamStreamForced + // directly, independent of IsStream. + if info.ChannelSetting.ForceUpstreamStream && !info.IsStream { + request.Stream = lo.ToPtr(true) + info.UpstreamStreamForced = true + // Inject stream_options.include_usage so the upstream returns actual + // usage in the final SSE chunk. Without this, the buffered handler + // falls back to estimated token counts, hurting billing accuracy. + if info.SupportStreamOptions && request.StreamOptions == nil { + request.StreamOptions = &dto.StreamOptions{ + IncludeUsage: true, + } + } + } + // Strip StreamOptions for channels that don't support them, but only + // when we did not inject it ourselves via ForceUpstreamStream. The + // forced-stream path (above) injects IncludeUsage for billing accuracy; + // nil-ing it here would make that injection dead code for any channel + // whose type is not OpenAI/Azure (e.g. DeepSeek with SupportStreamOptions). + // However, when the channel does not support StreamOptions at all + // (SupportStreamOptions=false), we must still strip them — even in + // forced-stream mode — to avoid sending unsupported fields upstream. + // + // Ordering dependency: shouldPreserveStreamOptions reads + // info.UpstreamStreamForced which is set inside the ForceUpstreamStream + // block above (line ~261). This guard MUST stay below that block. + shouldPreserveStreamOptions := info.UpstreamStreamForced && info.SupportStreamOptions + if !shouldPreserveStreamOptions && + info.ChannelType != constant.ChannelTypeOpenAI && + info.ChannelType != constant.ChannelTypeAzure { request.StreamOptions = nil } if info.ChannelType == constant.ChannelTypeOpenRouter { @@ -366,14 +423,17 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn return request, nil } +// ConvertRerankRequest converts a rerank request to the upstream request body. func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) { return request, nil } +// ConvertEmbeddingRequest converts an embedding request to the upstream request body. func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) { return request, nil } +// ConvertAudioRequest converts an audio request to the upstream request body. func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) { a.ResponseFormat = request.ResponseFormat if info.RelayMode == relayconstant.RelayModeAudioSpeech { @@ -440,6 +500,7 @@ func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInf } } +// ConvertImageRequest converts an image generation request to the upstream request body. func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) { switch info.RelayMode { case relayconstant.RelayModeImagesEdits: @@ -601,6 +662,7 @@ func detectImageMimeType(filename string) string { } } +// ConvertOpenAIResponsesRequest converts an OpenAI Responses API request to the upstream request body. func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) { // 转换模型推理力度后缀 effort, originModel := reasoning.ParseOpenAIReasoningEffortFromModelSuffix(request.Model) @@ -620,6 +682,7 @@ func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommo return request, nil } +// DoRequest executes the upstream HTTP request and returns the raw response. func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) { if info.RelayMode == relayconstant.RelayModeAudioTranscription || info.RelayMode == relayconstant.RelayModeAudioTranslation || @@ -632,6 +695,10 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request } } +// DoResponse dispatches the upstream HTTP response to the appropriate handler +// based on relay mode and stream state. When UpstreamStreamForced is true, it +// routes to OaiBufferedStreamHandler to aggregate the upstream SSE into a +// single JSON response for the non-streaming client. func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) { switch info.RelayMode { case relayconstant.RelayModeRealtime: @@ -659,7 +726,11 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom case relayconstant.RelayModeResponsesCompact: usage, err = OaiResponsesCompactionHandler(c, resp) default: - if info.IsStream { + if info.UpstreamStreamForced { + // Forced upstream stream: the upstream returned SSE but the client + // asked for non-streaming. Aggregate into a single JSON response. + usage, err = OaiBufferedStreamHandler(c, info, resp) + } else if info.IsStream { usage, err = OaiStreamHandler(c, info, resp) } else { usage, err = OpenaiHandler(c, info, resp) @@ -668,6 +739,7 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom return } +// GetModelList returns the list of models configured for this channel. func (a *Adaptor) GetModelList() []string { switch a.ChannelType { case constant.ChannelType360: @@ -685,6 +757,7 @@ func (a *Adaptor) GetModelList() []string { } } +// GetChannelName returns the human-readable channel type name. func (a *Adaptor) GetChannelName() string { switch a.ChannelType { case constant.ChannelType360: diff --git a/relay/channel/openai/adaptor_test.go b/relay/channel/openai/adaptor_test.go new file mode 100644 index 000000000000..82cfa7171361 --- /dev/null +++ b/relay/channel/openai/adaptor_test.go @@ -0,0 +1,249 @@ +package openai + +import ( + "bytes" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/relaykit/types" + "github.com/gin-gonic/gin" + "github.com/samber/lo" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestConvertOpenAIRequest_ForceUpstreamStream(t *testing.T) { + tests := []struct { + name string + clientStream *bool + forceUpstream bool + supportStreamOpts bool + channelType int + wantStreamSent bool // what the upstream should receive + wantForcedFlag bool // whether UpstreamStreamForced should be set + wantIsStream bool // info.IsStream after conversion (always matches client request) + wantStreamOptions bool // whether StreamOptions.IncludeUsage should be true + }{ + { + name: "client non-stream + force -> upstream stream + forced flag", + clientStream: lo.ToPtr(false), + forceUpstream: true, + supportStreamOpts: true, + wantStreamSent: true, + wantForcedFlag: true, + wantIsStream: false, // IsStream reflects client request, not forced upstream + wantStreamOptions: true, + }, + { + name: "client stream + force -> upstream stream, no forced flag", + clientStream: lo.ToPtr(true), + forceUpstream: true, + supportStreamOpts: true, + wantStreamSent: true, + wantForcedFlag: false, + wantIsStream: true, + wantStreamOptions: false, // forced flag not set, so StreamOptions not injected by force path + }, + { + name: "client non-stream + no force -> upstream non-stream, no forced flag", + clientStream: lo.ToPtr(false), + forceUpstream: false, + supportStreamOpts: true, + wantStreamSent: false, + wantForcedFlag: false, + wantIsStream: false, + wantStreamOptions: false, + }, + { + name: "force + no stream options support -> stream injected but no StreamOptions", + clientStream: lo.ToPtr(false), + forceUpstream: true, + supportStreamOpts: false, + wantStreamSent: true, + wantForcedFlag: true, + wantIsStream: false, + wantStreamOptions: false, + }, + { + // Bug-injection: non-OpenAI/Azure channel with force + stream options + // support. The buggy code unconditionally nils StreamOptions for + // non-OpenAI/Azure channels, making the IncludeUsage injection + // dead code. The fix scopes the nil-out with + // !info.UpstreamStreamForced so the forced-stream path keeps its + // StreamOptions. + name: "force + non-OpenAI channel + stream options support -> StreamOptions preserved", + clientStream: lo.ToPtr(false), + forceUpstream: true, + supportStreamOpts: true, + channelType: constant.ChannelTypeDeepSeek, + wantStreamSent: true, + wantForcedFlag: true, + wantIsStream: false, + wantStreamOptions: true, + }, + { + // force + non-OpenAI channel + NO stream options support -> StreamOptions stripped. + // Even in forced-stream mode, if the channel doesn't support + // StreamOptions, we must nil them to avoid upstream 400 errors. + name: "force + non-OpenAI channel + no stream options support -> StreamOptions stripped", + clientStream: lo.ToPtr(false), + forceUpstream: true, + supportStreamOpts: false, + channelType: constant.ChannelTypeDeepSeek, + wantStreamSent: true, + wantForcedFlag: true, + wantIsStream: false, + wantStreamOptions: false, + }, + { + // nil *bool clientStream should be treated as false (non-stream), + // matching lo.FromPtrOr's default. Force should still apply. + name: "nil clientStream + force -> upstream stream + forced flag", + clientStream: nil, + forceUpstream: true, + supportStreamOpts: true, + channelType: constant.ChannelTypeOpenAI, + wantStreamSent: true, + wantForcedFlag: true, + wantIsStream: false, + wantStreamOptions: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + + chType := tt.channelType + if chType == 0 { + chType = constant.ChannelTypeOpenAI + } + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelType: chType, + UpstreamModelName: "test-model", + ChannelSetting: dto.ChannelSettings{ForceUpstreamStream: tt.forceUpstream}, + SupportStreamOptions: tt.supportStreamOpts, + }, + RelayFormat: types.RelayFormatOpenAI, + } + + request := &dto.GeneralOpenAIRequest{ + Model: "test-model", + Stream: tt.clientStream, + } + + adaptor := &Adaptor{ChannelType: chType} + result, err := adaptor.ConvertOpenAIRequest(c, info, request) + require.NoError(t, err) + + returnedRequest, ok := result.(*dto.GeneralOpenAIRequest) + require.True(t, ok, "expected *GeneralOpenAIRequest, got %T", result) + + assert.Equal(t, tt.wantStreamSent, lo.FromPtrOr(returnedRequest.Stream, false), + "upstream stream field mismatch") + assert.Equal(t, tt.wantForcedFlag, info.UpstreamStreamForced, + "UpstreamStreamForced flag mismatch") + assert.Equal(t, tt.wantIsStream, info.IsStream, + "info.IsStream must match the original client request, not the forced upstream stream") + + if tt.wantStreamOptions { + require.NotNil(t, returnedRequest.StreamOptions, + "StreamOptions should be injected when stream is forced and provider supports it") + assert.True(t, returnedRequest.StreamOptions.IncludeUsage, + "StreamOptions.IncludeUsage must be true") + } else { + assert.Nil(t, returnedRequest.StreamOptions, + "StreamOptions must not be injected when stream is not forced or provider does not support it") + } + }) + } +} + +func TestDoResponse_RoutesForcedStreamToBufferedHandler(t *testing.T) { + gin.SetMode(gin.TestMode) + + // Set a valid streaming timeout to avoid NewTicker panic in OaiStreamHandler + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + // SSE response that OaiBufferedStreamHandler can aggregate + sseBody := strings.Join([]string{ + `data: {"id":"chatcmpl-x","object":"chat.completion.chunk","created":1,"model":"test","choices":[{"index":0,"delta":{"role":"assistant","content":"Hi"},"finish_reason":"stop"}]}`, + `data: [DONE]`, + ``, + }, "\n") + + tests := []struct { + name string + upstreamStreamForced bool + wantJSON bool // true = buffered handler (JSON), false = stream handler (SSE) + }{ + { + name: "forced stream -> buffered handler (JSON response)", + upstreamStreamForced: true, + wantJSON: true, + }, + { + name: "normal stream -> stream handler (SSE response)", + upstreamStreamForced: false, + wantJSON: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resp := &http.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseBody))), + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + c.Set(common.RequestIdKey, "test-req") + + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelType: constant.ChannelTypeOpenAI, + UpstreamModelName: "test", + }, + IsStream: true, + UpstreamStreamForced: tt.upstreamStreamForced, + RelayFormat: types.RelayFormatOpenAI, + } + + adaptor := &Adaptor{ChannelType: constant.ChannelTypeOpenAI} + usage, apiErr := adaptor.DoResponse(c, resp, info) + require.Nil(t, apiErr) + require.NotNil(t, usage) + + contentType := w.Header().Get("Content-Type") + body := w.Body.String() + if tt.wantJSON { + // Buffered handler produces a single JSON object with + // Content-Type application/json. + assert.Contains(t, contentType, "application/json", + "forced stream route must return application/json") + assert.Contains(t, body, "chat.completion", + "expected JSON response from buffered handler") + assert.NotContains(t, body, "data: ", + "buffered handler should not produce SSE data: lines") + } else { + // Stream handler writes SSE chunks with "data:" prefix + assert.True(t, strings.Contains(body, "data:") || strings.Contains(contentType, "text/event-stream"), + "expected SSE response from stream handler, got: %s", body[:min(100, len(body))]) + } + }) + } +} diff --git a/relay/channel/openai/buffered_stream.go b/relay/channel/openai/buffered_stream.go new file mode 100644 index 000000000000..135a40a7475b --- /dev/null +++ b/relay/channel/openai/buffered_stream.go @@ -0,0 +1,260 @@ +package openai + +import ( + "bufio" + "fmt" + "net/http" + "sort" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/logger" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relay/helper" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/relaykit/types" + "github.com/QuantumNous/new-api/service" + + "github.com/gin-gonic/gin" +) + +// OaiBufferedStreamHandler reads an upstream SSE stream of chat.completion.chunk +// events, aggregates them into a single chat.completion JSON, and writes it to +// the client. Used when ForceUpstreamStream is enabled and the client requested +// non-streaming. +func OaiBufferedStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { + if resp == nil || resp.Body == nil { + return nil, types.NewOpenAIError(fmt.Errorf("invalid response"), types.ErrorCodeBadResponse, http.StatusInternalServerError) + } + defer service.CloseResponseBodyGracefully(resp) + + var ( + accumulatedContent = make(map[int]string) // per choice index + accumulatedReasoning = make(map[int]string) // per choice index + accumulatedToolCalls = make(map[int]map[int]*dto.ToolCallResponse) // choiceIdx -> tcIdx -> tc + finishReason = make(map[int]string) // per choice index + model = info.UpstreamModelName + responseId = helper.GetResponseID(c) + created = time.Now().Unix() + usage *dto.Usage + ) + + scanner := helper.NewStreamScanner(resp.Body) + scanner.Split(bufio.ScanLines) + for scanner.Scan() { + line := scanner.Text() + if len(line) < 6 || line[:5] != "data:" { + continue + } + data := strings.TrimSpace(line[5:]) + if data == "[DONE]" { + break + } + if data == "" { + continue // heartbeat / keep-alive + } + + // Check for upstream error event before parsing as stream response. + // Cheap pre-check avoids double-unmarshal on normal chunks. + if strings.Contains(data, "\"error\"") { + var simpleResp dto.SimpleResponse + if err := common.UnmarshalJsonStr(data, &simpleResp); err == nil && simpleResp.Error != nil { + apiErr := simpleResp.GetOpenAIError() + if apiErr != nil { + return nil, types.NewOpenAIError(fmt.Errorf("upstream error: %s", apiErr.Message), types.ErrorCodeBadResponse, http.StatusBadGateway) + } + return nil, types.NewOpenAIError(fmt.Errorf("upstream returned error event"), types.ErrorCodeBadResponse, http.StatusBadGateway) + } + } + + var streamResp dto.ChatCompletionsStreamResponse + if err := common.UnmarshalJsonStr(data, &streamResp); err != nil { + logger.LogError(c, "failed to unmarshal buffered stream chunk: "+err.Error()) + continue + } + + if streamResp.Usage != nil { + usage = streamResp.Usage + } + if model == "" && streamResp.Model != "" { + model = streamResp.Model + } + if len(streamResp.Choices) > 0 { + for _, choice := range streamResp.Choices { + idx := choice.Index + if choice.Delta.GetContentString() != "" { + accumulatedContent[idx] += choice.Delta.GetContentString() + } + if choice.Delta.GetReasoningContent() != "" { + accumulatedReasoning[idx] += choice.Delta.GetReasoningContent() + } + if len(choice.Delta.ToolCalls) > 0 { + if accumulatedToolCalls[idx] == nil { + accumulatedToolCalls[idx] = make(map[int]*dto.ToolCallResponse) + } + for _, tc := range choice.Delta.ToolCalls { + tcIdx := 0 + if tc.Index != nil { + tcIdx = *tc.Index + } + if existing, ok := accumulatedToolCalls[idx][tcIdx]; !ok { + tcCopy := tc + accumulatedToolCalls[idx][tcIdx] = &tcCopy + } else { + existing.Function.Arguments += tc.Function.Arguments + if tc.Function.Name != "" { + existing.Function.Name = tc.Function.Name + } + if tc.ID != "" { + existing.ID = tc.ID + } + } + } + } + if choice.FinishReason != nil && *choice.FinishReason != "" { + finishReason[idx] = *choice.FinishReason + } + } + } + } + + if err := scanner.Err(); err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError) + } + + // Determine all choice indices that received content, reasoning, + // tool calls, or a finish reason. Tool-call-only choices (no content, + // no reasoning, no finish_reason) must not be silently dropped. + allIndices := make(map[int]bool) + for idx := range accumulatedContent { + allIndices[idx] = true + } + for idx := range accumulatedReasoning { + allIndices[idx] = true + } + for idx := range accumulatedToolCalls { + allIndices[idx] = true + } + for idx := range finishReason { + allIndices[idx] = true + } + + // Build choices for all indices (sorted for deterministic output) + var sortedIndices []int + for idx := range allIndices { + sortedIndices = append(sortedIndices, idx) + } + if len(sortedIndices) == 0 { + sortedIndices = []int{0} + } + sort.Ints(sortedIndices) + + var choices []dto.OpenAITextResponseChoice + for _, idx := range sortedIndices { + fr := finishReason[idx] + if fr == "" { + fr = constant.FinishReasonStop + } + choice := dto.OpenAITextResponseChoice{ + Index: idx, + FinishReason: fr, + } + choice.Message.Role = "assistant" + choice.Message.Content = accumulatedContent[idx] + if accumulatedReasoning[idx] != "" { + rc := accumulatedReasoning[idx] + choice.Message.ReasoningContent = &rc + } + if tcMap, ok := accumulatedToolCalls[idx]; ok && len(tcMap) > 0 { + var tcs []dto.ToolCallResponse + var tcKeys []int + for k := range tcMap { + tcKeys = append(tcKeys, k) + } + sort.Ints(tcKeys) + for _, k := range tcKeys { + tcs = append(tcs, *tcMap[k]) + } + choice.Message.SetToolCalls(tcs) + } + choices = append(choices, choice) + } + + // Usage fallback: aggregate all content across choices for estimation. + // Include reasoning content and tool-call arguments so the estimate + // matches what ProcessStreamResponse would compute for the same stream. + if usage == nil || usage.TotalTokens == 0 { + totalContent := "" + for _, idx := range sortedIndices { + totalContent += accumulatedContent[idx] + accumulatedReasoning[idx] + if tcMap, ok := accumulatedToolCalls[idx]; ok { + tcKeys := make([]int, 0, len(tcMap)) + for k := range tcMap { + tcKeys = append(tcKeys, k) + } + sort.Ints(tcKeys) + for _, k := range tcKeys { + totalContent += tcMap[k].Function.Name + tcMap[k].Function.Arguments + } + } + } + usage = service.ResponseText2Usage(c, totalContent, info.UpstreamModelName, info.GetEstimatePromptTokens()) + } + // Guard against nil usage from the fallback estimator. If the upstream + // returned no usage object AND the estimator returned nil (e.g. empty + // model name or zero content), dereferencing *usage below would panic. + if usage == nil { + logger.LogWarn(c, "buffered stream: usage estimator returned nil, using zero usage") + usage = &dto.Usage{} + } + + textResponse := dto.OpenAITextResponse{ + Id: responseId, + Object: "chat.completion", + Created: created, + Model: model, + Choices: choices, + Usage: *usage, + } + + responseBody, err := common.Marshal(textResponse) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeJsonMarshalFailed, http.StatusInternalServerError) + } + + // Apply channel-specific usage post-processing (e.g. DeepSeek cache-hit + // token migration) and re-marshal if usage changed. Matches the pattern + // in OpenaiHandler (P2-3). + applyUsagePostProcessing(info, &textResponse.Usage, responseBody) + if textResponse.Usage.PromptTokensDetails.CachedTokens != usage.PromptTokensDetails.CachedTokens { + responseBody, err = common.Marshal(textResponse) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeJsonMarshalFailed, http.StatusInternalServerError) + } + usage = &textResponse.Usage + } + + // Count billable tool calls for special tool pricing, matching + // OaiStreamHandler and OpenaiHandler (P2-3). Iterate per choice. + for _, tcMap := range accumulatedToolCalls { + for _, tc := range tcMap { + if tc.Function.Name != "" { + info.CountBillableToolCall(dto.BuildInCallFunctionCall, tc.Function.Name) + } + } + } + + // The buffered handler has fully parsed and rebuilt the response as a + // single JSON object. Write it directly with the correct Content-Type + // instead of using IOCopyBytesGracefully, which would copy the upstream's + // text/event-stream header and mislead strict clients (P0-1). + c.Writer.Header().Set("Content-Type", "application/json") + c.Writer.Header().Set("Content-Length", fmt.Sprintf("%d", len(responseBody))) + c.Writer.WriteHeader(http.StatusOK) + _, _ = c.Writer.Write(responseBody) + c.Writer.Flush() + + return usage, nil +} diff --git a/relay/channel/openai/buffered_stream_test.go b/relay/channel/openai/buffered_stream_test.go new file mode 100644 index 000000000000..0157c69a424a --- /dev/null +++ b/relay/channel/openai/buffered_stream_test.go @@ -0,0 +1,471 @@ +package openai + +import ( + "bytes" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestOaiBufferedStreamHandler_AggregatesContent(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,"model":"kimi-k2.6","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,"model":"kimi-k2.6","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,"model":"kimi-k2.6","choices":[{"index":0,"delta":{"content":"!"},"finish_reason":"stop"}]}`, + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,"model":"kimi-k2.6","choices":[],"usage":{"prompt_tokens":10,"completion_tokens":3,"total_tokens":13}}`, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseBody))), + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "kimi-k2.6"}, + IsStream: true, + UpstreamStreamForced: true, + } + + usage, apiErr := OaiBufferedStreamHandler(c, info, resp) + require.Nil(t, apiErr) + require.NotNil(t, usage) + assert.Equal(t, 3, usage.CompletionTokens) + assert.Equal(t, 10, usage.PromptTokens) + + body := w.Body.String() + assert.Contains(t, body, `"object":"chat.completion"`) + assert.Contains(t, body, "Hello world!") +} + +func TestOaiBufferedStreamHandler_AggregatesReasoningContent(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: {"id":"chatcmpl-2","object":"chat.completion.chunk","created":1,"model":"deepseek-r1","choices":[{"index":0,"delta":{"role":"assistant","reasoning_content":"Thinking"},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-2","object":"chat.completion.chunk","created":1,"model":"deepseek-r1","choices":[{"index":0,"delta":{"reasoning_content":" about it"},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-2","object":"chat.completion.chunk","created":1,"model":"deepseek-r1","choices":[{"index":0,"delta":{"content":"Answer"},"finish_reason":"stop"}]}`, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseBody))), + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "deepseek-r1"}, + IsStream: true, + UpstreamStreamForced: true, + } + + usage, apiErr := OaiBufferedStreamHandler(c, info, resp) + require.Nil(t, apiErr) + require.NotNil(t, usage) + + body := w.Body.String() + assert.Contains(t, body, "Answer") + // Verify reasoning_content is present in the response + var textResp dto.OpenAITextResponse + require.NoError(t, common.Unmarshal([]byte(body), &textResp)) + require.Len(t, textResp.Choices, 1) + assert.NotEmpty(t, textResp.Choices[0].Message.GetReasoningContent()) + assert.Contains(t, textResp.Choices[0].Message.GetReasoningContent(), "Thinking about it") +} + +func TestOaiBufferedStreamHandler_AggregatesToolCalls(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: {"id":"chatcmpl-3","object":"chat.completion.chunk","created":1,"model":"gpt-4","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":""}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-3","object":"chat.completion.chunk","created":1,"model":"gpt-4","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"loc"}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-3","object":"chat.completion.chunk","created":1,"model":"gpt-4","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"ation\":\"NYC\"}"}}]},"finish_reason":"tool_calls"}]}`, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseBody))), + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-4"}, + IsStream: true, + UpstreamStreamForced: true, + } + + usage, apiErr := OaiBufferedStreamHandler(c, info, resp) + require.Nil(t, apiErr) + require.NotNil(t, usage) + + body := w.Body.String() + var textResp dto.OpenAITextResponse + require.NoError(t, common.Unmarshal([]byte(body), &textResp)) + require.Len(t, textResp.Choices, 1) + toolCalls := textResp.Choices[0].Message.ParseToolCalls() + require.Len(t, toolCalls, 1) + assert.Equal(t, "get_weather", toolCalls[0].Function.Name) + assert.Equal(t, `{"location":"NYC"}`, toolCalls[0].Function.Arguments) + assert.Equal(t, "tool_calls", textResp.Choices[0].FinishReason) +} + +func TestOaiBufferedStreamHandler_MissingFinishChunk(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: {"id":"chatcmpl-4","object":"chat.completion.chunk","created":1,"model":"test","choices":[{"index":0,"delta":{"content":"Hi"},"finish_reason":null}]}`, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseBody))), + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "test"}, + IsStream: true, + UpstreamStreamForced: true, + } + + usage, apiErr := OaiBufferedStreamHandler(c, info, resp) + require.Nil(t, apiErr) + require.NotNil(t, usage) + + body := w.Body.String() + assert.Contains(t, body, "Hi") + assert.Contains(t, body, `"object":"chat.completion"`) +} + +// TestOaiBufferedStreamHandler_ToolCallBilling verifies that the buffered +// handler counts billable tool calls for special tool pricing, matching +// OaiStreamHandler and OpenaiHandler (P2-3). Without this, forced streams +// that return tool_calls skip per-call tool billing. +func TestOaiBufferedStreamHandler_ToolCallBilling(t *testing.T) { + gin.SetMode(gin.TestMode) + + operation_setting.SetToolPriceForTest("my_priced_fn", 5.0) + t.Cleanup(func() { + operation_setting.DeleteToolPriceForTest("my_priced_fn") + }) + + sseBody := strings.Join([]string{ + `data: {"id":"chatcmpl-tb","object":"chat.completion.chunk","created":1,"model":"gpt-4","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"my_priced_fn","arguments":""}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-tb","object":"chat.completion.chunk","created":1,"model":"gpt-4","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{}"}}]},"finish_reason":"tool_calls"}]}`, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseBody))), + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-4"}, + OriginModelName: "gpt-4", + IsStream: true, + UpstreamStreamForced: true, + } + + _, apiErr := OaiBufferedStreamHandler(c, info, resp) + require.Nil(t, apiErr) + + require.NotNil(t, info.ResponsesUsageInfo, "ResponsesUsageInfo must be initialized by CountBillableToolCall") + require.Contains(t, info.ResponsesUsageInfo.BuiltInTools, "my_priced_fn", + "priced tool call must be counted for billing") + assert.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools["my_priced_fn"].CallCount, + "call count must be 1 for a single tool invocation") +} + +// TestOaiBufferedStreamHandler_UsagePostProcessing verifies that the buffered +// handler applies channel-specific usage post-processing (e.g. DeepSeek +// cache-hit token migration), matching OpenaiHandler (P2-3). Without this, +// DeepSeek cached-token billing is silently lost on forced streams. +func TestOaiBufferedStreamHandler_UsagePostProcessing(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: {"id":"chatcmpl-up","object":"chat.completion.chunk","created":1,"model":"deepseek-chat","choices":[{"index":0,"delta":{"content":"Hi"},"finish_reason":"stop"}]}`, + `data: {"id":"chatcmpl-up","object":"chat.completion.chunk","created":1,"model":"deepseek-chat","choices":[],"usage":{"prompt_tokens":10,"completion_tokens":1,"total_tokens":11,"prompt_cache_hit_tokens":5}}`, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseBody))), + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelType: constant.ChannelTypeDeepSeek, + UpstreamModelName: "deepseek-chat", + }, + OriginModelName: "deepseek-chat", + IsStream: true, + UpstreamStreamForced: true, + } + + usage, apiErr := OaiBufferedStreamHandler(c, info, resp) + require.Nil(t, apiErr) + require.NotNil(t, usage) + + assert.Equal(t, 5, usage.PromptTokensDetails.CachedTokens, + "DeepSeek prompt_cache_hit_tokens must be migrated to PromptTokensDetails.CachedTokens by applyUsagePostProcessing") + + body := w.Body.String() + var textResp dto.OpenAITextResponse + require.NoError(t, common.Unmarshal([]byte(body), &textResp)) + assert.Equal(t, 5, textResp.Usage.PromptTokensDetails.CachedTokens, + "response body must reflect migrated cached tokens") +} + +// TestOaiBufferedStreamHandler_ContentTypeIsJSON verifies that the buffered +// handler sets Content-Type to application/json, not the upstream's +// text/event-stream. A strict HTTP client rejects a JSON body declared as +// text/event-stream (P0-1). +func TestOaiBufferedStreamHandler_ContentTypeIsJSON(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: {"id":"chatcmpl-ct","object":"chat.completion.chunk","created":1,"model":"test","choices":[{"index":0,"delta":{"content":"Hi"},"finish_reason":"stop"}]}`, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseBody))), + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "test"}, + IsStream: true, + UpstreamStreamForced: true, + } + + _, apiErr := OaiBufferedStreamHandler(c, info, resp) + require.Nil(t, apiErr) + + contentType := w.Header().Get("Content-Type") + assert.Contains(t, contentType, "application/json", + "Content-Type must be application/json, got: %s", contentType) + assert.NotContains(t, contentType, "text/event-stream", + "Content-Type must not leak upstream text/event-stream") +} + +// TestOaiBufferedStreamHandler_UpstreamErrorEvent verifies that an error event +// in the SSE stream is surfaced as an API error -- the handler returns a +// NewAPIError and nil usage so the client sees the failure and billing is not +// charged for an empty success. +func TestOaiBufferedStreamHandler_UpstreamErrorEvent(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: {"error":{"message":"rate limited","type":"rate_limit_error"}}`, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseBody))), + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "test"}, + IsStream: true, + UpstreamStreamForced: true, + } + + usage, apiErr := OaiBufferedStreamHandler(c, info, resp) + // An upstream error event is a real error, not data. The handler must + // return a NewAPIError and nil usage so the client sees the failure + // and billing is not charged for an empty success. + require.NotNil(t, apiErr, "handler must return API error for in-stream error event") + assert.Nil(t, usage, "usage must be nil when upstream returns error") + assert.Contains(t, apiErr.Error(), "upstream error", "error message must mention upstream") +} + +// TestOaiBufferedStreamHandler_MalformedDataLines verifies that malformed +// data lines in the SSE stream are skipped without causing errors, and that +// empty data payloads (heartbeats) are skipped without ending aggregation. +func TestOaiBufferedStreamHandler_MalformedDataLines(t *testing.T) { + gin.SetMode(gin.TestMode) + + sseBody := strings.Join([]string{ + `data: not-json`, + `data: `, // heartbeat — empty payload, should be skipped + `data: {"id":"x","object":"chat.completion.chunk","created":1,"model":"test","choices":[{"index":0,"delta":{"content":"OK"},"finish_reason":"stop"}]}`, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseBody))), + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "test"}, + IsStream: true, + UpstreamStreamForced: true, + } + + _, apiErr := OaiBufferedStreamHandler(c, info, resp) + assert.Nil(t, apiErr, "handler must skip malformed lines without error") + body := w.Body.String() + assert.Contains(t, body, "OK", "valid content after malformed lines must be aggregated") +} + +// TestOaiBufferedStreamHandler_ToolCallOnlyChoiceNotDropped verifies that a +// choice index which receives only tool_calls (no content, no reasoning, no +// finish_reason) is still present in the aggregated response. Without +// collecting indices from accumulatedToolCalls, such a choice is silently +// dropped from allIndices and never appears in the output. +func TestOaiBufferedStreamHandler_ToolCallOnlyChoiceNotDropped(t *testing.T) { + gin.SetMode(gin.TestMode) + + // Choice index 1 receives ONLY tool_calls — no content, no finish_reason. + // The buggy code only collected indices from content/reasoning/finishReason, + // so index 1 would be dropped. The fix adds accumulatedToolCalls to the + // index collection. + sseBody := strings.Join([]string{ + `data: {"id":"chatcmpl-tc-only","object":"chat.completion.chunk","created":1,"model":"gpt-4","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":"stop"}]}`, + `data: {"id":"chatcmpl-tc-only","object":"chat.completion.chunk","created":1,"model":"gpt-4","choices":[{"index":1,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_x","type":"function","function":{"name":"do_thing","arguments":"{}"}}]}}]}`, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseBody))), + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-4"}, + IsStream: true, + UpstreamStreamForced: true, + } + + _, apiErr := OaiBufferedStreamHandler(c, info, resp) + require.Nil(t, apiErr) + + body := w.Body.String() + var textResp dto.OpenAITextResponse + require.NoError(t, common.Unmarshal([]byte(body), &textResp)) + require.Len(t, textResp.Choices, 2, "both choice indices must appear: index 0 (content) and index 1 (tool_calls only)") + + // Find choice with index 1 + var choice1 *dto.OpenAITextResponseChoice + for i := range textResp.Choices { + if textResp.Choices[i].Index == 1 { + choice1 = &textResp.Choices[i] + break + } + } + require.NotNil(t, choice1, "choice index 1 (tool_calls only) must not be dropped") + toolCalls := choice1.Message.ParseToolCalls() + require.Len(t, toolCalls, 1) + assert.Equal(t, "do_thing", toolCalls[0].Function.Name) +} + +// TestOaiBufferedStreamHandler_NilUsageNoPanic verifies that the handler does +// not panic when the upstream returns no usage object and the fallback +// estimator returns nil (simulated via empty content + empty model name). +// Without the nil guard, `Usage: *usage` dereferences a nil pointer. +func TestOaiBufferedStreamHandler_NilUsageNoPanic(t *testing.T) { + gin.SetMode(gin.TestMode) + + // SSE stream with no usage chunk and no content (so ResponseText2Usage + // gets empty string). An empty model name makes the estimator return nil. + sseBody := strings.Join([]string{ + `data: {"id":"chatcmpl-nil","object":"chat.completion.chunk","created":1,"model":"","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: 200, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sseBody))), + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequestWithContext(t.Context(), "POST", "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: ""}, + IsStream: true, + UpstreamStreamForced: true, + } + + // Must not panic + require.NotPanics(t, func() { + usage, apiErr := OaiBufferedStreamHandler(c, info, resp) + require.Nil(t, apiErr) + require.NotNil(t, usage, "usage must be non-nil even when estimator returns nil") + }) + + body := w.Body.String() + assert.Contains(t, body, "chat.completion", "response must still be valid JSON") +} + diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index b0bb19bdca3b..87c814c22c72 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -23,6 +23,7 @@ import ( "github.com/tidwall/gjson" ) +// ThinkingContentInfo tracks thinking/reasoning content state during relay processing. type ThinkingContentInfo struct { IsFirstThinkingContent bool SendLastThinkingContent bool @@ -40,21 +41,25 @@ const ( // host code and adaptors compiling unchanged. type ClaudeConvertInfo = convmeta.ClaudeConvertInfo +// RerankerInfo holds parameters for reranker requests. type RerankerInfo struct { Documents []any ReturnDocuments bool } +// BuildInToolInfo holds built-in tool configuration for the relay. type BuildInToolInfo struct { ToolName string CallCount int SearchContextSize string } +// ResponsesUsageInfo tracks usage statistics for OpenAI Responses API requests. type ResponsesUsageInfo struct { BuiltInTools map[string]*BuildInToolInfo } +// ChannelMeta holds channel-level metadata used across the relay pipeline. type ChannelMeta struct { ChannelType int ChannelId int @@ -75,11 +80,13 @@ type ChannelMeta struct { SupportStreamOptions bool // 是否支持流式选项 } +// TokenCountMeta tracks token counting state for billing and rate limiting. type TokenCountMeta struct { //promptTokens int estimatePromptTokens int } +// RelayInfo is the central context object passed through the relay pipeline, carrying request metadata, channel settings, and per-attempt state. type RelayInfo struct { TokenId int TokenKey string @@ -93,6 +100,7 @@ type RelayInfo struct { isFirstResponse bool //SendLastReasoningResponse bool IsStream bool + UpstreamStreamForced bool // true when client requested non-stream but upstream was forced to stream IsGeminiBatchEmbedding bool IsPlayground bool UsePrice bool @@ -185,6 +193,7 @@ type RelayInfo struct { *TaskRelayInfo } +// InitChannelMeta initializes channel metadata from the gin context and channel configuration. func (info *RelayInfo) InitChannelMeta(c *gin.Context) { channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType) paramOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelParamOverride) @@ -247,6 +256,7 @@ func (info *RelayInfo) InitChannelMeta(c *gin.Context) { } } +// ToString returns a JSON representation of RelayInfo for debugging. func (info *RelayInfo) ToString() string { if info == nil { return "RelayInfo" @@ -637,6 +647,7 @@ func GenRelayInfo(c *gin.Context, relayFormat types.RelayFormat, request dto.Req return info, nil } +// InitRequestConversionChain resets the request conversion format chain. func (info *RelayInfo) InitRequestConversionChain() { if info == nil { return @@ -650,6 +661,7 @@ func (info *RelayInfo) InitRequestConversionChain() { info.RequestConversionChain = []types.RelayFormat{info.RelayFormat} } +// AppendRequestConversion appends a relay format to the conversion chain. func (info *RelayInfo) AppendRequestConversion(format types.RelayFormat) { if info == nil { return @@ -668,6 +680,7 @@ func (info *RelayInfo) AppendRequestConversion(format types.RelayFormat) { info.RequestConversionChain = append(info.RequestConversionChain, format) } +// GetFinalRequestRelayFormat returns the last format in the request conversion chain. func (info *RelayInfo) GetFinalRequestRelayFormat() types.RelayFormat { if info == nil { return "" @@ -718,6 +731,7 @@ func (info *RelayInfo) SetEstimatePromptTokens(promptTokens int) { info.estimatePromptTokens = promptTokens } +// GetEstimatePromptTokens returns the estimated prompt token count for fallback usage calculation. func (info *RelayInfo) GetEstimatePromptTokens() int { if info == nil { return 0 @@ -732,6 +746,7 @@ func (info *RelayInfo) GetEstimatePromptTokens() int { var _ convmeta.Meta = (*RelayInfo)(nil) +// GetOriginModelName returns the model name as specified by the client. func (info *RelayInfo) GetOriginModelName() string { if info == nil { return "" @@ -739,6 +754,7 @@ func (info *RelayInfo) GetOriginModelName() string { return info.OriginModelName } +// GetUpstreamModelName returns the model name sent to the upstream. func (info *RelayInfo) GetUpstreamModelName() string { if info == nil || info.ChannelMeta == nil { return "" @@ -746,8 +762,10 @@ func (info *RelayInfo) GetUpstreamModelName() string { return info.UpstreamModelName } +// HasChannelMeta returns true if the RelayInfo has non-nil channel metadata. func (info *RelayInfo) HasChannelMeta() bool { return info != nil && info.ChannelMeta != nil } +// GetChannelID returns the numeric channel identifier. func (info *RelayInfo) GetChannelID() int { if info == nil || info.ChannelMeta == nil { return 0 @@ -755,6 +773,7 @@ func (info *RelayInfo) GetChannelID() int { return info.ChannelId } +// GetChannelType returns the channel type constant. func (info *RelayInfo) GetChannelType() int { if info == nil || info.ChannelMeta == nil { return 0 @@ -762,10 +781,12 @@ func (info *RelayInfo) GetChannelType() int { return info.ChannelType } +// GetIsStream returns whether the current request is a streaming request. func (info *RelayInfo) GetIsStream() bool { return info != nil && info.IsStream } +// GetReasoningEffort returns the reasoning effort level for the request. func (info *RelayInfo) GetReasoningEffort() string { if info == nil { return "" @@ -773,6 +794,7 @@ func (info *RelayInfo) GetReasoningEffort() string { return info.ReasoningEffort } +// SetReasoningEffort sets the reasoning effort level for the request. func (info *RelayInfo) SetReasoningEffort(effort string) { if info == nil { return @@ -780,6 +802,7 @@ func (info *RelayInfo) SetReasoningEffort(effort string) { info.ReasoningEffort = strings.TrimSpace(effort) } +// EnsureClaudeConvertInfo returns the Claude conversion metadata, initializing it if needed. func (info *RelayInfo) EnsureClaudeConvertInfo() *convmeta.ClaudeConvertInfo { if info == nil { return &convmeta.ClaudeConvertInfo{ @@ -794,6 +817,7 @@ func (info *RelayInfo) EnsureClaudeConvertInfo() *convmeta.ClaudeConvertInfo { return info.ClaudeConvertInfo } +// GetSendResponseCount returns the number of response chunks sent to the client. func (info *RelayInfo) GetSendResponseCount() int { if info == nil { return 0 @@ -801,6 +825,7 @@ func (info *RelayInfo) GetSendResponseCount() int { return info.SendResponseCount } +// IncrSendResponseCount increments the response chunk counter. func (info *RelayInfo) IncrSendResponseCount() { if info == nil { return @@ -839,6 +864,7 @@ func (info *RelayInfo) ConvOptions() *convmeta.Options { return options } +// SetFirstResponseTime records the time of the first response byte. func (info *RelayInfo) SetFirstResponseTime() { if info.isFirstResponse { info.FirstResponseTime = time.Now() @@ -846,10 +872,12 @@ func (info *RelayInfo) SetFirstResponseTime() { } } +// HasSendResponse returns true if at least one response chunk has been sent. func (info *RelayInfo) HasSendResponse() bool { return info.FirstResponseTime.After(info.StartTime) } +// TaskRelayInfo holds metadata for async task relay requests. type TaskRelayInfo struct { Action string OriginTaskID string @@ -865,6 +893,7 @@ type TaskRelayInfo struct { LockedChannel any } +// TaskSubmitReq represents the submission payload for async task requests. type TaskSubmitReq struct { Prompt string `json:"prompt"` Model string `json:"model,omitempty"` @@ -878,14 +907,17 @@ type TaskSubmitReq struct { Metadata map[string]interface{} `json:"metadata,omitempty"` } +// GetPrompt returns the text prompt from the task submission. func (t *TaskSubmitReq) GetPrompt() string { return t.Prompt } +// HasImage returns true if the task submission contains image content. func (t *TaskSubmitReq) HasImage() bool { return len(t.Images) > 0 } +// UnmarshalJSON implements custom JSON unmarshalling for TaskSubmitReq. func (t *TaskSubmitReq) UnmarshalJSON(data []byte) error { type Alias TaskSubmitReq aux := &struct { @@ -932,6 +964,7 @@ func (t *TaskSubmitReq) UnmarshalJSON(data []byte) error { return nil } +// UnmarshalMetadata deserializes the task metadata into the given struct. func (t *TaskSubmitReq) UnmarshalMetadata(v any) error { metadata := t.Metadata if metadata != nil { @@ -947,6 +980,7 @@ func (t *TaskSubmitReq) UnmarshalMetadata(v any) error { return nil } +// TaskInfo holds task status and result information for async requests. type TaskInfo struct { Code int `json:"code"` TaskID string `json:"task_id"` diff --git a/relay/common/relay_info_test.go b/relay/common/relay_info_test.go index 42a0f8567bfe..eaf6c8b03f7f 100644 --- a/relay/common/relay_info_test.go +++ b/relay/common/relay_info_test.go @@ -176,3 +176,11 @@ func TestInitChannelMetaRestoresRequestReasoningEffortForRetry(t *testing.T) { info.InitChannelMeta(ctx) assert.Equal(t, "max", info.ReasoningEffort) } + +func TestUpstreamStreamForcedField(t *testing.T) { + info := &RelayInfo{} + info.UpstreamStreamForced = true + if !info.UpstreamStreamForced { + t.Error("UpstreamStreamForced field not settable") + } +} diff --git a/relaykit/dto/channel_settings.go b/relaykit/dto/channel_settings.go index 4b4e71911283..895516a4092b 100644 --- a/relaykit/dto/channel_settings.go +++ b/relaykit/dto/channel_settings.go @@ -10,6 +10,7 @@ import ( "github.com/QuantumNous/new-api/relaykit/types" ) +// ChannelSettings holds per-channel configuration for relay behavior, authentication, and routing. type ChannelSettings struct { ForceFormat bool `json:"force_format,omitempty"` ThinkingToContent bool `json:"thinking_to_content,omitempty"` @@ -23,6 +24,13 @@ type ChannelSettings struct { // HTTP2ConnectionShards spreads HTTP/2 traffic across N independent transports // (1-8). Zero/unset means 1. Ignored when HTTPProtocol is "http1". HTTP2ConnectionShards int `json:"http2_connection_shards,omitempty"` + // ForceUpstreamStream makes new-api send stream=true to the upstream even + // when the downstream client requested non-streaming. The SSE response is + // aggregated server-side into a single JSON. Mutually exclusive with + // PassThroughBodyEnabled. Note: StreamOptions.IncludeUsage is injected + // only for OpenAI and Azure channels (SupportStreamOptions=true); other + // OpenAI-compatible channels (e.g. DeepSeek) will use estimated usage. + ForceUpstreamStream bool `json:"force_upstream_stream,omitempty"` } const ( @@ -51,6 +59,19 @@ func (s *ChannelSettings) ValidateHTTPTransport() error { return nil } +// ValidateForceUpstreamStream rejects configurations where ForceUpstreamStream +// and PassThroughBodyEnabled are both enabled, since they are mutually exclusive. +func (s *ChannelSettings) ValidateForceUpstreamStream() error { + if s == nil { + return nil + } + if s.ForceUpstreamStream && s.PassThroughBodyEnabled { + return fmt.Errorf("force_upstream_stream and pass_through_body_enabled are mutually exclusive") + } + return nil +} + +// VertexKeyType identifies the authentication method for Google Vertex AI channels. type VertexKeyType string const ( @@ -58,6 +79,7 @@ const ( VertexKeyTypeAPIKey VertexKeyType = "api_key" ) +// AwsKeyType identifies the authentication method for AWS Bedrock channels. type AwsKeyType string const ( @@ -65,6 +87,7 @@ const ( AwsKeyTypeApiKey AwsKeyType = "api_key" ) +// ChannelOtherSettings holds supplementary channel configuration not covered by ChannelSettings. type ChannelOtherSettings struct { AzureResponsesVersion string `json:"azure_responses_version,omitempty"` VertexKeyType VertexKeyType `json:"vertex_key_type,omitempty"` // "json" or "api_key" @@ -87,6 +110,7 @@ type ChannelOtherSettings struct { AdvancedCustom *AdvancedCustomConfig `json:"advanced_custom,omitempty"` } +// IsOpenRouterEnterprise returns true if the channel uses OpenRouter enterprise routing. func (s *ChannelOtherSettings) IsOpenRouterEnterprise() bool { if s == nil || s.OpenRouterEnterprise == nil { return false @@ -111,10 +135,12 @@ const ( AdvancedCustomAuthTypeQuery = "query" ) +// AdvancedCustomConfig holds advanced per-model routing and endpoint configuration. type AdvancedCustomConfig struct { Routes []AdvancedCustomRoute `json:"advanced_routes,omitempty"` } +// AdvancedCustomRoute defines a custom routing rule for a specific model or pattern. type AdvancedCustomRoute struct { IncomingPath string `json:"incoming_path,omitempty"` UpstreamPath string `json:"upstream_path,omitempty"` @@ -123,6 +149,7 @@ type AdvancedCustomRoute struct { Auth *AdvancedCustomRouteAuth `json:"auth,omitempty"` } +// AdvancedCustomRouteAuth holds authentication overrides for a custom route. type AdvancedCustomRouteAuth struct { Type string `json:"type,omitempty"` Name string `json:"name,omitempty"` @@ -222,6 +249,7 @@ func (c *AdvancedCustomConfig) SupportsPathForModel(requestPath string, model st return ok } +// SupportedEndpointTypesForModel returns the endpoint types supported by the custom config for the given model. func (c *AdvancedCustomConfig) SupportedEndpointTypesForModel(model string) []types.EndpointType { if c == nil { return nil @@ -367,6 +395,7 @@ func IsAdvancedCustomConverterAllowed(converter string) bool { } } +// Validate checks the advanced custom configuration for internal consistency. func (c *AdvancedCustomConfig) Validate() error { if c == nil { return fmt.Errorf("advanced_custom is required") diff --git a/relaykit/dto/channel_settings_test.go b/relaykit/dto/channel_settings_test.go index e84988731bf8..3eb03b274196 100644 --- a/relaykit/dto/channel_settings_test.go +++ b/relaykit/dto/channel_settings_test.go @@ -642,3 +642,39 @@ func TestChannelSettingsValidateHTTPTransport(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "http2_connection_shards") } + +func TestValidateForceUpstreamStream(t *testing.T) { + tests := []struct { + name string + s ChannelSettings + wantErr bool + }{ + { + name: "force alone is ok", + s: ChannelSettings{ForceUpstreamStream: true}, + wantErr: false, + }, + { + name: "passthrough alone is ok", + s: ChannelSettings{PassThroughBodyEnabled: true}, + wantErr: false, + }, + { + name: "both enabled is rejected", + s: ChannelSettings{ForceUpstreamStream: true, PassThroughBodyEnabled: true}, + wantErr: true, + }, + { + name: "neither is ok", + s: ChannelSettings{}, + wantErr: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if err := tt.s.ValidateForceUpstreamStream(); (err != nil) != tt.wantErr { + t.Errorf("ValidateForceUpstreamStream() error = %v, wantErr %v", err, tt.wantErr) + } + }) + } +}