From 269e4ff390594ef532a6587c48c6966fa617ce8e Mon Sep 17 00:00:00 2001 From: CaIon Date: Sat, 11 Jul 2026 13:13:56 +0800 Subject: [PATCH] feat(image): enhance image stream handling with client disconnect logic and billing adjustments --- relay/channel/openai/image_stream_test.go | 280 +++++++++++++++++++++- relay/channel/openai/relay_image.go | 145 +++++++---- 2 files changed, 377 insertions(+), 48 deletions(-) diff --git a/relay/channel/openai/image_stream_test.go b/relay/channel/openai/image_stream_test.go index a9b1e0b2..1adde5c9 100644 --- a/relay/channel/openai/image_stream_test.go +++ b/relay/channel/openai/image_stream_test.go @@ -1,14 +1,17 @@ package openai import ( + "context" "io" "net/http" "net/http/httptest" "strings" + "sync" "testing" "github.com/QuantumNous/new-api/constant" relaycommon "github.com/QuantumNous/new-api/relay/common" + relayconstant "github.com/QuantumNous/new-api/relay/constant" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -32,6 +35,37 @@ func newImageTestContext(t *testing.T, body, contentType string, isStream bool) return c, recorder, resp, info } +func TestOpenaiImageDoResponseUsesInfoIsStream(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + body := `{"created":1710000000,"data":[{"b64_json":"image"}]}` + + t.Run("non-stream response stays JSON", func(t *testing.T) { + c, recorder, resp, info := newImageTestContext(t, body, "application/json", false) + info.RelayMode = relayconstant.RelayModeImagesGenerations + + usage, err := (&Adaptor{}).DoResponse(c, resp, info) + + require.Nil(t, err) + require.NotNil(t, usage) + require.Equal(t, body, recorder.Body.String()) + }) + + t.Run("stream response converts JSON to SSE", func(t *testing.T) { + c, recorder, resp, info := newImageTestContext(t, body, "application/json", true) + info.RelayMode = relayconstant.RelayModeImagesGenerations + + usage, err := (&Adaptor{}).DoResponse(c, resp, info) + + require.Nil(t, err) + require.NotNil(t, usage) + require.Contains(t, recorder.Body.String(), `event: image_generation.completed`) + require.Contains(t, recorder.Body.String(), `data: [DONE]`) + }) +} + // TestOpenaiImageStreamHandlerForwardsSSEAndUsage covers the core SSE path: // chunks are forwarded with rebuilt event lines, usage is extracted and // normalized (input_tokens -> prompt_tokens with details), and [DONE] is @@ -56,6 +90,8 @@ func TestOpenaiImageStreamHandlerForwardsSSEAndUsage(t *testing.T) { }, "\n") c, recorder, resp, info := newImageTestContext(t, body, "text/event-stream", true) + info.PriceData.UsePrice = true + info.PriceData.AddOtherRatio("n", 3) usage, err := OpenaiImageStreamHandler(c, info, resp) require.Nil(t, err) @@ -69,6 +105,181 @@ func TestOpenaiImageStreamHandlerForwardsSSEAndUsage(t *testing.T) { require.Contains(t, recorder.Body.String(), `data: {"usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7,"input_tokens_details":{"image_tokens":2,"text_tokens":1}}}`) require.Contains(t, recorder.Body.String(), `data: [DONE]`) require.Equal(t, "text/event-stream", recorder.Header().Get("Content-Type")) + require.Equal(t, 3.0, info.PriceData.OtherRatios()["n"], "streams without completed events keep the requested count") +} + +func TestOpenaiImageStreamHandlerUsesCompletedEventCount(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := strings.Join([]string{ + `data: {"type":"image_generation.partial_image","partial_image_index":0,"b64_json":"partial"}`, + ``, + `data: {"type":"image_generation.completed","b64_json":"first"}`, + ``, + `data: {"type":"image_edit.completed","b64_json":"second","usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7}}`, + ``, + `data: [DONE]`, + ``, + }, "\n") + + c, _, resp, info := newImageTestContext(t, body, "text/event-stream", true) + info.PriceData.UsePrice = true + info.PriceData.AddOtherRatio("n", 3) + + usage, err := OpenaiImageStreamHandler(c, info, resp) + + require.Nil(t, err) + require.Equal(t, 7, usage.TotalTokens) + require.Equal(t, 2.0, info.PriceData.OtherRatios()["n"]) +} + +// blockingBody serves one SSE chunk, then blocks until Close (the scanner's +// cleanup) and returns EOF — keeping the upstream "open" while the client-side +// disconnect is simulated elsewhere. +type blockingBody struct { + mu sync.Mutex + sent bool + chunk []byte + closed chan struct{} +} + +func (b *blockingBody) Read(p []byte) (int, error) { + b.mu.Lock() + if !b.sent { + b.sent = true + n := copy(p, b.chunk) + b.mu.Unlock() + return n, nil + } + b.mu.Unlock() + <-b.closed + return 0, io.EOF +} + +func (b *blockingBody) Close() error { + b.mu.Lock() + defer b.mu.Unlock() + select { + case <-b.closed: + default: + close(b.closed) + } + return nil +} + +// cancelAfterWriter cancels the request context right after the payload +// containing needle has been written to the client, simulating a client that +// disconnects after receiving that event. Cancelling from the write side (not +// the upstream read side) makes the abort deterministic: the handler has +// already processed and counted the event when the disconnect fires. +type cancelAfterWriter struct { + gin.ResponseWriter + needle string + cancel context.CancelFunc + once sync.Once +} + +func (w *cancelAfterWriter) Write(p []byte) (int, error) { + n, err := w.ResponseWriter.Write(p) + if strings.Contains(string(p), w.needle) { + w.once.Do(w.cancel) + } + return n, err +} + +func (w *cancelAfterWriter) WriteString(s string) (int, error) { + n, err := io.WriteString(w.ResponseWriter, s) + if strings.Contains(s, w.needle) { + w.once.Do(w.cancel) + } + return n, err +} + +func newDisconnectingImageStream(t *testing.T, sseBody, disconnectAfter string) (*gin.Context, *httptest.ResponseRecorder, *http.Response, *relaycommon.RelayInfo) { + t.Helper() + c, recorder, resp, info := newImageTestContext(t, "", "text/event-stream", true) + ctx, cancel := context.WithCancel(c.Request.Context()) + t.Cleanup(cancel) + c.Request = c.Request.WithContext(ctx) + c.Writer = &cancelAfterWriter{ResponseWriter: c.Writer, needle: disconnectAfter, cancel: cancel} + resp.Body = &blockingBody{ + chunk: []byte(sseBody), + closed: make(chan struct{}), + } + return c, recorder, resp, info +} + +// TestOpenaiImageStreamHandlerClientDisconnectKeepsRequestedCount guards the +// billing invariant: completed-event counting must not lower the charge when +// the client aborts the stream. Upstream already generated (and charged for) +// all requested images, so a disconnect after the first completed event keeps +// the requested n instead of dropping it to 1. +func TestOpenaiImageStreamHandlerClientDisconnectKeepsRequestedCount(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := "data: {\"type\":\"image_generation.completed\",\"b64_json\":\"first\"}\n\n" + c, recorder, resp, info := newDisconnectingImageStream(t, body, "first") + info.PriceData.UsePrice = true + info.PriceData.AddOtherRatio("n", 3) + + usage, err := OpenaiImageStreamHandler(c, info, resp) + + require.Nil(t, err) + require.NotNil(t, usage) + require.NotNil(t, info.StreamStatus) + // A client abort surfaces as client_gone (main-loop ctx watch) or + // handler_stop (failed client write); both must be treated as untrusted. + require.Contains(t, + []relaycommon.StreamEndReason{relaycommon.StreamEndReasonClientGone, relaycommon.StreamEndReasonHandlerStop}, + info.StreamStatus.EndReason) + require.Contains(t, recorder.Body.String(), `"b64_json":"first"`) + require.Equal(t, 3.0, info.PriceData.OtherRatios()["n"], "client abort must not reduce the billed image count") +} + +// TestOpenaiImageStreamHandlerClientDisconnectRaisesCount covers the other +// direction of the abort guard: when completed events already exceed the +// recorded n, the higher actual count is billed even though the client aborted. +func TestOpenaiImageStreamHandlerClientDisconnectRaisesCount(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := strings.Join([]string{ + `data: {"type":"image_generation.completed","b64_json":"first"}`, + ``, + `data: {"type":"image_generation.completed","b64_json":"second"}`, + ``, + ``, + }, "\n") + c, _, resp, info := newDisconnectingImageStream(t, body, "second") + info.PriceData.UsePrice = true + info.PriceData.AddOtherRatio("n", 1) + + usage, err := OpenaiImageStreamHandler(c, info, resp) + + require.Nil(t, err) + require.NotNil(t, usage) + require.NotNil(t, info.StreamStatus) + require.Contains(t, + []relaycommon.StreamEndReason{relaycommon.StreamEndReasonClientGone, relaycommon.StreamEndReasonHandlerStop}, + info.StreamStatus.EndReason) + require.Equal(t, 2.0, info.PriceData.OtherRatios()["n"], "completed events beyond the recorded n must raise the charge even on abort") } // TestOpenaiImageStreamHandlerWrapsJSONResponse covers the non-SSE fallback: @@ -78,9 +289,11 @@ func TestOpenaiImageStreamHandlerWrapsJSONResponse(t *testing.T) { gin.SetMode(gin.TestMode) t.Cleanup(func() { gin.SetMode(oldMode) }) - body := `{"created":1710000000,"data":[{"b64_json":"final","revised_prompt":"draw a cat"}],"usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7,"input_tokens_details":{"image_tokens":2,"text_tokens":1}}}` + body := `{"created":1710000000,"data":[{"b64_json":"first","revised_prompt":"draw a cat"},{"b64_json":"second"}],"usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7,"input_tokens_details":{"image_tokens":2,"text_tokens":1}}}` c, recorder, resp, info := newImageTestContext(t, body, "application/json", true) + info.PriceData.UsePrice = true + info.PriceData.AddOtherRatio("n", 3) usage, err := OpenaiImageStreamHandler(c, info, resp) require.Nil(t, err) @@ -93,9 +306,59 @@ func TestOpenaiImageStreamHandlerWrapsJSONResponse(t *testing.T) { require.Empty(t, recorder.Header().Get("Content-Length")) require.Contains(t, recorder.Body.String(), `event: image_generation.completed`) require.Contains(t, recorder.Body.String(), `"type":"image_generation.completed"`) - require.Contains(t, recorder.Body.String(), `"b64_json":"final"`) + require.Contains(t, recorder.Body.String(), `"b64_json":"first"`) + require.Contains(t, recorder.Body.String(), `"b64_json":"second"`) require.Contains(t, recorder.Body.String(), `"revised_prompt":"draw a cat"`) require.Contains(t, recorder.Body.String(), `data: [DONE]`) + require.Equal(t, 2, strings.Count(recorder.Body.String(), `event: image_generation.completed`)) + require.Equal(t, 2.0, info.PriceData.OtherRatios()["n"]) +} + +func TestOpenaiImageHandlerUsesPositiveActualCountForFixedPrice(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + longImage := strings.Repeat("a", 4096) + + tests := []struct { + name string + body string + usePrice bool + wantCount float64 + }{ + { + name: "fixed price uses data length", + body: `{"data":[{"b64_json":"` + longImage + `"},{"b64_json":"second"}]}`, + usePrice: true, + wantCount: 2, + }, + { + name: "empty data keeps requested count", + body: `{"data":[]}`, + usePrice: true, + wantCount: 3, + }, + { + name: "ratio billing ignores data length", + body: `{"data":[{"b64_json":"first"},{"b64_json":"second"}]}`, + usePrice: false, + wantCount: 3, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c, recorder, resp, info := newImageTestContext(t, tt.body, "application/json", false) + info.PriceData.UsePrice = tt.usePrice + info.PriceData.AddOtherRatio("n", 3) + + _, err := OpenaiImageHandler(c, info, resp) + + require.Nil(t, err) + require.Equal(t, tt.wantCount, info.PriceData.OtherRatios()["n"]) + require.Equal(t, tt.body, recorder.Body.String()) + }) + } } // TestOpenaiImageHandlersReturnJSONError covers JSON error responses for both @@ -132,6 +395,19 @@ func TestOpenaiImageHandlersReturnJSONError(t *testing.T) { require.Equal(t, "content moderation failed", err.ToOpenAIError().Message) require.Empty(t, recorder.Body.String()) }) + + t.Run("stream handler non-2xx stays JSON error", func(t *testing.T) { + c, recorder, resp, info := newImageTestContext(t, body, "application/json", true) + resp.StatusCode = http.StatusBadGateway + + usage, err := OpenaiImageStreamHandler(c, info, resp) + require.Nil(t, usage) + require.NotNil(t, err) + require.Equal(t, http.StatusBadGateway, err.StatusCode) + require.Equal(t, "content moderation failed", err.ToOpenAIError().Message) + require.Empty(t, recorder.Body.String()) + require.NotContains(t, recorder.Header().Get("Content-Type"), "text/event-stream") + }) } // TestOpenaiImageStreamHandlerRecordsUpstreamErrorEvent verifies that an error diff --git a/relay/channel/openai/relay_image.go b/relay/channel/openai/relay_image.go index 1eccbe2c..90fb8794 100644 --- a/relay/channel/openai/relay_image.go +++ b/relay/channel/openai/relay_image.go @@ -5,6 +5,7 @@ import ( "fmt" "io" "net/http" + "strconv" "strings" "time" @@ -17,8 +18,17 @@ import ( "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" ) +func updateOpenAIImageCount(info *relaycommon.RelayInfo, count int64) { + if info == nil || !info.PriceData.UsePrice || count <= 0 || count > int64(dto.MaxImageN) { + return + } + info.PriceData.AddOtherRatio("n", float64(count)) +} + // OpenaiImageHandler handles non-streaming OpenAI image responses // (generations/edits), returning the parsed usage for billing. func OpenaiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { @@ -39,6 +49,8 @@ func OpenaiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. return nil, types.WithOpenAIError(*oaiError, resp.StatusCode) } + updateOpenAIImageCount(info, gjson.GetBytes(responseBody, "data.#").Int()) + // 写入新的 response body service.IOCopyBytesGracefully(c, resp, responseBody) @@ -88,7 +100,7 @@ func OpenaiImageStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp return OpenaiImageHandler(c, info, resp) } if !strings.Contains(contentType, "text/event-stream") { - return OpenaiImageJSONAsStreamHandler(c, info, resp) + return openaiImageJSONAsStreamHandler(c, info, resp) } // Reuse the shared streaming engine (helper.StreamScannerHandler) so the // image streaming path gets the same ping keepalive, streaming-timeout @@ -98,6 +110,7 @@ func OpenaiImageStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp // field (real OpenAI image events keep event == type). usage := &dto.Usage{} var lastStreamData []byte + var completedImages int64 helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) { raw := common.StringToByteSlice(data) @@ -107,39 +120,64 @@ func OpenaiImageStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp // EndReason. HasErrors() flags the failure for logging/handling. sr.Error(fmt.Errorf("%s", extractOpenAIImageStreamErrorMessage(raw))) } - var usageResp dto.SimpleResponse - if err := common.Unmarshal(raw, &usageResp); err == nil { - normalizeOpenAIUsage(&usageResp.Usage) - if service.ValidUsage(&usageResp.Usage) { - usage = &usageResp.Usage + var chunk struct { + Type string `json:"type"` + Usage dto.Usage `json:"usage"` + } + if err := common.Unmarshal(raw, &chunk); err == nil { + normalizeOpenAIUsage(&chunk.Usage) + if service.ValidUsage(&chunk.Usage) { + usage = &chunk.Usage + } + if chunk.Type == "image_generation.completed" || chunk.Type == "image_edit.completed" { + completedImages++ } } - writeOpenaiImageStreamChunk(c, raw) + if err := writeOpenaiImageStreamChunk(c, raw); err != nil { + sr.Stop(err) + } }) // StreamScannerHandler consumes the upstream [DONE]; re-emit it so the // client still receives a terminal data: [DONE]. - if info != nil && info.StreamStatus != nil && info.StreamStatus.EndReason == relaycommon.StreamEndReasonDone { + if info.StreamStatus != nil && info.StreamStatus.EndReason == relaycommon.StreamEndReasonDone { helper.Done(c) } applyUsagePostProcessing(info, usage, lastStreamData) + // Only trust completedImages when upstream finished the stream (done/eof). + // On client-side aborts (client_gone, or handler_stop from a failed client + // write) the counter undercounts what upstream actually generated and + // charged, so keep the requested n — otherwise a client could pay for one + // image by disconnecting right after the first completed event. The abort + // guard only blocks lowering the charge: if completed events already + // exceed the recorded n, bill the higher actual count regardless. + if info.StreamStatus != nil { + upstreamFinished := info.StreamStatus.EndReason == relaycommon.StreamEndReasonDone || + info.StreamStatus.EndReason == relaycommon.StreamEndReasonEOF + requestedN := 1.0 + if n, ok := info.PriceData.OtherRatios()["n"]; ok { + requestedN = n + } + if upstreamFinished || float64(completedImages) > requestedN { + updateOpenAIImageCount(info, completedImages) + } + } return usage, nil } // writeOpenaiImageStreamChunk rebuilds the SSE frame for an image stream chunk: // it emits an "event:" line derived from the JSON "type" field (when present) // followed by the verbatim "data:" payload, mirroring helper.ResponseChunkData. -func writeOpenaiImageStreamChunk(c *gin.Context, data []byte) { +func writeOpenaiImageStreamChunk(c *gin.Context, data []byte) error { var payload struct { Type string `json:"type"` } _ = common.Unmarshal(data, &payload) if eventName := strings.TrimSpace(payload.Type); eventName != "" { - _ = helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: eventName}, string(data)) - return + return helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: eventName}, string(data)) } - _ = helper.StringData(c, string(data)) + return helper.StringData(c, string(data)) } // isOpenAIImageStreamErrorEvent detects upstream error chunks by JSON content @@ -192,7 +230,7 @@ func extractOpenAIImageStreamErrorMessage(data []byte) string { return "upstream image stream returned error event" } -func OpenaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { +func openaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { defer service.CloseResponseBodyGracefully(resp) responseBody, err := io.ReadAll(resp.Body) @@ -200,49 +238,75 @@ func OpenaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError) } - var imageResp dto.ImageResponse - if err := common.Unmarshal(responseBody, &imageResp); err != nil { + // Only decode usage/error. Do not Unmarshal data[] into dto.ImageResponse — + // b64_json values are large and would be copied into Go strings then + // re-marshaled for each SSE event. + var usageResp dto.SimpleResponse + if err := common.Unmarshal(responseBody, &usageResp); err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } - - var usageResp dto.SimpleResponse - _ = common.Unmarshal(responseBody, &usageResp) if oaiError := usageResp.GetOpenAIError(); oaiError != nil && oaiError.Type != "" { return nil, types.WithOpenAIError(*oaiError, resp.StatusCode) } normalizeOpenAIUsage(&usageResp.Usage) applyUsagePostProcessing(info, &usageResp.Usage, responseBody) + imageCount := gjson.GetBytes(responseBody, "data.#").Int() + updateOpenAIImageCount(info, imageCount) + helper.SetEventStreamHeaders(c) c.Status(http.StatusOK) - created := imageResp.Created + created := gjson.GetBytes(responseBody, "created").Int() if created == 0 { created = time.Now().Unix() } if info != nil { info.SetFirstResponseTime() } - for _, image := range imageResp.Data { - payload := map[string]any{ - "type": "image_generation.completed", - "created_at": created, + + validUsage := service.ValidUsage(&usageResp.Usage) + var usageJSON []byte + if validUsage { + usageJSON, err = common.Marshal(usageResp.Usage) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } - if image.Url != "" { - payload["url"] = image.Url + } + + for i := int64(0); i < imageCount; i++ { + image := gjson.GetBytes(responseBody, "data."+strconv.FormatInt(i, 10)) + payload := []byte(`{"type":"image_generation.completed"}`) + payload, err = sjson.SetBytes(payload, "created_at", created) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } - if image.B64Json != "" { - payload["b64_json"] = image.B64Json + if validUsage { + payload, err = sjson.SetRawBytes(payload, "usage", usageJSON) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + } } - if image.RevisedPrompt != "" { - payload["revised_prompt"] = image.RevisedPrompt + // b64_json goes last: every sjson.Set* reallocates the whole payload, + // so inserting the large blob after all small fields avoids re-copying + // multi-MB buffers. + for _, field := range []string{"url", "revised_prompt", "b64_json"} { + value := image.Get(field) + if value.Type != gjson.String || value.Raw == `""` { + continue + } + raw := []byte(value.Raw) + if value.Index > 0 { + raw = responseBody[value.Index : value.Index+len(value.Raw)] + } + payload, err = sjson.SetRawBytes(payload, field, raw) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + } } - if service.ValidUsage(&usageResp.Usage) { - payload["usage"] = usageResp.Usage - } - if err := writeOpenaiImageStreamPayload(c, "image_generation.completed", payload); err != nil { + if writeErr := helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: "image_generation.completed"}, string(payload)); writeErr != nil { if info != nil && info.StreamStatus != nil { - info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, err) + info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, writeErr) } return &usageResp.Usage, nil } @@ -254,7 +318,7 @@ func OpenaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, return &usageResp.Usage, nil } if info != nil { - info.ReceivedResponseCount += len(imageResp.Data) + info.ReceivedResponseCount += int(imageCount) if info.StreamStatus == nil { info.StreamStatus = relaycommon.NewStreamStatus() } @@ -263,17 +327,6 @@ func OpenaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, return &usageResp.Usage, nil } -func writeOpenaiImageStreamPayload(c *gin.Context, eventName string, payload any) error { - data, err := common.Marshal(payload) - if err != nil { - return err - } - if eventName != "" { - return helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: eventName}, string(data)) - } - return helper.StringData(c, string(data)) -} - func writeOpenaiImageStreamDone(c *gin.Context) error { return helper.StringData(c, "[DONE]") }