chaos/chaos-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(openai): align image streaming relay governance

Browse Source 
Route OpenAI image streaming through shared stream handling, split image/realtime/usage helpers for maintainability, and include the related image request and rate limit updates.
This commit is contained in:
CaIon
2026-06-10 17:47:37 +08:00
parent 867d8acfc3
commit 59a93cf5c7
14 changed files with 845 additions and 979 deletions
Download Patch File Download Diff File Expand all files Collapse all files
relay/channel/openai/image_edit_test.go
+65 -88
View File
@@ -16,106 +16,83 @@ import (
"github.com/stretchr/testify/require"
)
// TestConvertImageEditRequestKeepsValidMultipartStreamFields verifies multipart replay.
func TestConvertImageEditRequestKeepsValidMultipartStreamFields(t *testing.T) {
// TestConvertImageEditRequestMultipart verifies that ConvertImageRequest
// re-serializes multipart image edit requests with all fields (including
// stream) and the file intact, both when the form was already parsed and when
// it must be re-parsed from the reusable body.
func TestConvertImageEditRequestMultipart(t *testing.T) {
gin.SetMode(gin.TestMode)
var body bytes.Buffer
writer := multipart.NewWriter(&body)
require.NoError(t, writer.WriteField("model", "gpt-image-1"))
require.NoError(t, writer.WriteField("prompt", "edit this image"))
require.NoError(t, writer.WriteField("stream", "true"))
require.NoError(t, writer.WriteField("partial_images", "3"))
part, err := writer.CreateFormFile("image", "input.png")
require.NoError(t, err)
_, err = part.Write([]byte("fake image"))
require.NoError(t, err)
require.NoError(t, writer.Close())
newMultipartContext := func(t *testing.T, prompt string) *gin.Context {
var body bytes.Buffer
writer := multipart.NewWriter(&body)
require.NoError(t, writer.WriteField("model", "gpt-image-1"))
require.NoError(t, writer.WriteField("prompt", prompt))
require.NoError(t, writer.WriteField("stream", "true"))
require.NoError(t, writer.WriteField("partial_images", "3"))
part, err := writer.CreateFormFile("image", "input.png")
require.NoError(t, err)
_, err = part.Write([]byte("fake image"))
require.NoError(t, err)
require.NoError(t, writer.Close())
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
require.NoError(t, c.Request.ParseMultipartForm(32<<20))
info := &relaycommon.RelayInfo{
RelayMode: relayconstant.RelayModeImagesEdits,
}
request := dto.ImageRequest{
Model: "gpt-image-1",
Prompt: "edit this image",
Stream: true,
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
return c
}
converted, err := (&Adaptor{}).ConvertImageRequest(c, info, request)
require.NoError(t, err)
convertAndReplay := func(t *testing.T, c *gin.Context, prompt string) {
info := &relaycommon.RelayInfo{
RelayMode: relayconstant.RelayModeImagesEdits,
}
request := dto.ImageRequest{
Model: "gpt-image-1",
Prompt: prompt,
Stream: common.GetPointer(true),
}
convertedBody, ok := converted.(*bytes.Buffer)
require.True(t, ok)
converted, err := (&Adaptor{}).ConvertImageRequest(c, info, request)
require.NoError(t, err)
convertedBody, ok := converted.(*bytes.Buffer)
require.True(t, ok)
contentType := c.Request.Header.Get("Content-Type")
replayedRequest := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(convertedBody.Bytes()))
replayedRequest.Header.Set("Content-Type", contentType)
require.NoError(t, replayedRequest.ParseMultipartForm(32<<20))
replayedRequest := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(convertedBody.Bytes()))
replayedRequest.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
require.NoError(t, replayedRequest.ParseMultipartForm(32<<20))
require.Equal(t, "gpt-image-1", replayedRequest.PostForm.Get("model"))
require.Equal(t, "edit this image", replayedRequest.PostForm.Get("prompt"))
require.Equal(t, "true", replayedRequest.PostForm.Get("stream"))
require.Equal(t, "3", replayedRequest.PostForm.Get("partial_images"))
require.Len(t, replayedRequest.MultipartForm.File["image"], 1)
require.Equal(t, "gpt-image-1", replayedRequest.PostForm.Get("model"))
require.Equal(t, prompt, replayedRequest.PostForm.Get("prompt"))
require.Equal(t, "true", replayedRequest.PostForm.Get("stream"))
require.Equal(t, "3", replayedRequest.PostForm.Get("partial_images"))
require.Len(t, replayedRequest.MultipartForm.File["image"], 1)
file, err := replayedRequest.MultipartForm.File["image"][0].Open()
require.NoError(t, err)
defer file.Close()
fileBytes, err := io.ReadAll(file)
require.NoError(t, err)
require.Equal(t, []byte("fake image"), fileBytes)
}
// TestConvertImageEditRequestParsesReusableMultipartWhenFormIsMissing verifies fallback parsing.
func TestConvertImageEditRequestParsesReusableMultipartWhenFormIsMissing(t *testing.T) {
gin.SetMode(gin.TestMode)
var body bytes.Buffer
writer := multipart.NewWriter(&body)
require.NoError(t, writer.WriteField("model", "gpt-image-1"))
require.NoError(t, writer.WriteField("prompt", "edit without pre-parsed form"))
require.NoError(t, writer.WriteField("stream", "true"))
part, err := writer.CreateFormFile("image", "input.png")
require.NoError(t, err)
_, err = part.Write([]byte("fake image"))
require.NoError(t, err)
require.NoError(t, writer.Close())
file, err := replayedRequest.MultipartForm.File["image"][0].Open()
require.NoError(t, err)
defer file.Close()
fileBytes, err := io.ReadAll(file)
require.NoError(t, err)
require.Equal(t, []byte("fake image"), fileBytes)
}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
t.Run("with pre-parsed form", func(t *testing.T) {
prompt := "edit this image"
c := newMultipartContext(t, prompt)
require.NoError(t, c.Request.ParseMultipartForm(32<<20))
storage, err := common.GetBodyStorage(c)
require.NoError(t, err)
c.Request.Body = io.NopCloser(storage)
c.Request.MultipartForm = nil
c.Request.PostForm = nil
convertAndReplay(t, c, prompt)
})
info := &relaycommon.RelayInfo{
RelayMode: relayconstant.RelayModeImagesEdits,
}
request := dto.ImageRequest{
Model: "gpt-image-1",
Prompt: "edit without pre-parsed form",
Stream: true,
}
t.Run("re-parses reusable body when form is missing", func(t *testing.T) {
prompt := "edit without pre-parsed form"
c := newMultipartContext(t, prompt)
converted, err := (&Adaptor{}).ConvertImageRequest(c, info, request)
require.NoError(t, err)
storage, err := common.GetBodyStorage(c)
require.NoError(t, err)
c.Request.Body = io.NopCloser(storage)
c.Request.MultipartForm = nil
c.Request.PostForm = nil
convertedBody, ok := converted.(*bytes.Buffer)
require.True(t, ok)
replayedRequest := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(convertedBody.Bytes()))
replayedRequest.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
require.NoError(t, replayedRequest.ParseMultipartForm(32<<20))
require.Equal(t, "edit without pre-parsed form", replayedRequest.PostForm.Get("prompt"))
require.Equal(t, "true", replayedRequest.PostForm.Get("stream"))
require.Len(t, replayedRequest.MultipartForm.File["image"], 1)
convertAndReplay(t, c, prompt)
})
}