feat(image): enhance image stream handling with client disconnect logic and billing adjustments

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

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]")
 }