mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-07 17:08:33 +08:00
fix(images): report decoded OAuth output dimensions
ChatGPT OAuth may normalize explicit image size controls to auto while returning smaller raster bytes. Decode final image headers so Images API metadata and tier billing reflect the actual output for both streaming and non-streaming OAuth responses.
This commit is contained in:
@@ -0,0 +1,92 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/base64"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"image"
|
||||
_ "image/jpeg"
|
||||
_ "image/png"
|
||||
"io"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const maxOpenAIImageDimensionProbeBytes int64 = 1 << 20
|
||||
|
||||
func detectOpenAIImageResultSize(encoded string) string {
|
||||
payload := strings.TrimSpace(encoded)
|
||||
if strings.HasPrefix(strings.ToLower(payload), "data:") {
|
||||
comma := strings.IndexByte(payload, ',')
|
||||
if comma < 0 || comma+1 >= len(payload) {
|
||||
return ""
|
||||
}
|
||||
payload = strings.TrimSpace(payload[comma+1:])
|
||||
}
|
||||
if payload == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
for _, encoding := range []*base64.Encoding{base64.StdEncoding, base64.RawStdEncoding} {
|
||||
decoded := base64.NewDecoder(encoding, strings.NewReader(payload))
|
||||
buffered := bufio.NewReader(io.LimitReader(decoded, maxOpenAIImageDimensionProbeBytes))
|
||||
prefix, _ := buffered.Peek(30)
|
||||
if width, height, ok := detectOpenAIWebPDimensions(prefix); ok {
|
||||
return fmt.Sprintf("%dx%d", width, height)
|
||||
}
|
||||
cfg, _, err := image.DecodeConfig(buffered)
|
||||
if err != nil || cfg.Width <= 0 || cfg.Height <= 0 {
|
||||
continue
|
||||
}
|
||||
return fmt.Sprintf("%dx%d", cfg.Width, cfg.Height)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func detectOpenAIWebPDimensions(header []byte) (int, int, bool) {
|
||||
if len(header) < 16 || string(header[:4]) != "RIFF" || string(header[8:12]) != "WEBP" {
|
||||
return 0, 0, false
|
||||
}
|
||||
|
||||
switch string(header[12:16]) {
|
||||
case "VP8X":
|
||||
if len(header) < 30 {
|
||||
return 0, 0, false
|
||||
}
|
||||
width := 1 + int(header[24]) + int(header[25])<<8 + int(header[26])<<16
|
||||
height := 1 + int(header[27]) + int(header[28])<<8 + int(header[29])<<16
|
||||
return width, height, width > 0 && height > 0
|
||||
case "VP8 ":
|
||||
if len(header) < 30 || string(header[23:26]) != "\x9d\x01\x2a" {
|
||||
return 0, 0, false
|
||||
}
|
||||
width := int(binary.LittleEndian.Uint16(header[26:28]) & 0x3fff)
|
||||
height := int(binary.LittleEndian.Uint16(header[28:30]) & 0x3fff)
|
||||
return width, height, width > 0 && height > 0
|
||||
case "VP8L":
|
||||
if len(header) < 25 || header[20] != 0x2f {
|
||||
return 0, 0, false
|
||||
}
|
||||
width := 1 + int(header[21]) + int(header[22]&0x3f)<<8
|
||||
height := 1 + int(header[22]>>6) + int(header[23])<<2 + int(header[24]&0x0f)<<10
|
||||
return width, height, width > 0 && height > 0
|
||||
default:
|
||||
return 0, 0, false
|
||||
}
|
||||
}
|
||||
|
||||
func reconcileOpenAIResponsesImageResultSizes(results []openAIResponsesImageResult, firstMeta *openAIResponsesImageResult) {
|
||||
for i := range results {
|
||||
// ChatGPT OAuth can normalize requested controls to "auto". The final
|
||||
// image bytes are authoritative for response metadata and tier billing.
|
||||
if actualSize := detectOpenAIImageResultSize(results[i].Result); actualSize != "" {
|
||||
results[i].Size = actualSize
|
||||
}
|
||||
}
|
||||
if firstMeta == nil || len(results) == 0 {
|
||||
return
|
||||
}
|
||||
if size := strings.TrimSpace(results[0].Size); size != "" {
|
||||
firstMeta.Size = size
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,174 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
"image/jpeg"
|
||||
"image/png"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestDetectOpenAIImageResultSize(t *testing.T) {
|
||||
pngEncoded := encodeOpenAIImageTestPNG(t, 1672, 941)
|
||||
jpegEncoded := encodeOpenAIImageTestJPEG(t, 640, 360)
|
||||
webpVP8XEncoded := encodeOpenAIImageTestWebPVP8X(1920, 1080)
|
||||
webpVP8Encoded := encodeOpenAIImageTestWebPVP8(1280, 720)
|
||||
webpVP8LEncoded := encodeOpenAIImageTestWebPVP8L(640, 480)
|
||||
|
||||
require.Equal(t, "1672x941", detectOpenAIImageResultSize(pngEncoded))
|
||||
require.Equal(t, "1672x941", detectOpenAIImageResultSize(strings.TrimRight(pngEncoded, "=")))
|
||||
require.Equal(t, "1672x941", detectOpenAIImageResultSize("data:image/png;base64,"+pngEncoded))
|
||||
require.Equal(t, "640x360", detectOpenAIImageResultSize(jpegEncoded))
|
||||
require.Equal(t, "1920x1080", detectOpenAIImageResultSize(webpVP8XEncoded))
|
||||
require.Equal(t, "1280x720", detectOpenAIImageResultSize(webpVP8Encoded))
|
||||
require.Equal(t, "640x480", detectOpenAIImageResultSize(webpVP8LEncoded))
|
||||
require.Empty(t, detectOpenAIImageResultSize("data:image/png;base64"))
|
||||
require.Empty(t, detectOpenAIImageResultSize("not-image-data"))
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForwardImages_OAuthUsesDecodedOutputDimensions(t *testing.T) {
|
||||
run := runOpenAIOAuthImageActualSizeTest(t, false)
|
||||
|
||||
require.Equal(t, "3840x2160", gjson.GetBytes(run.upstream.lastBody, "tools.0.size").String())
|
||||
require.Equal(t, "low", gjson.GetBytes(run.upstream.lastBody, "tools.0.quality").String())
|
||||
require.Equal(t, "1672x941", gjson.Get(run.recorder.Body.String(), "size").String())
|
||||
require.Equal(t, "auto", gjson.Get(run.recorder.Body.String(), "quality").String())
|
||||
require.Equal(t, []string{"1672x941"}, run.result.ImageOutputSizes)
|
||||
|
||||
ApplyOpenAIImageBillingResolution(run.result)
|
||||
require.Equal(t, ImageBillingSize2K, run.result.ImageSize)
|
||||
require.Equal(t, "1672x941", run.result.ImageOutputSize)
|
||||
require.Equal(t, ImageSizeSourceOutput, run.result.ImageSizeSource)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForwardImages_OAuthStreamingUsesDecodedOutputDimensions(t *testing.T) {
|
||||
run := runOpenAIOAuthImageActualSizeTest(t, true)
|
||||
|
||||
events := parseOpenAIImageTestSSEEvents(run.recorder.Body.String())
|
||||
completed, ok := findOpenAIImageTestSSEEvent(events, "image_generation.completed")
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "1672x941", gjson.Get(completed.Data, "size").String())
|
||||
require.Equal(t, "auto", gjson.Get(completed.Data, "quality").String())
|
||||
require.Equal(t, []string{"1672x941"}, run.result.ImageOutputSizes)
|
||||
}
|
||||
|
||||
type openAIOAuthImageActualSizeTestRun struct {
|
||||
result *OpenAIForwardResult
|
||||
recorder *httptest.ResponseRecorder
|
||||
upstream *httpUpstreamRecorder
|
||||
}
|
||||
|
||||
func runOpenAIOAuthImageActualSizeTest(t *testing.T, stream bool) openAIOAuthImageActualSizeTestRun {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
body := []byte(fmt.Sprintf(`{"model":"gpt-image-2","prompt":"draw a test chart","size":"3840x2160","quality":"low","output_format":"png","stream":%t}`, stream))
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = req
|
||||
c.Set("api_key", &APIKey{ID: 42})
|
||||
|
||||
encoded := encodeOpenAIImageTestPNG(t, 1672, 941)
|
||||
upstreamBody := fmt.Sprintf(
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"created_at\":1710000000,\"tools\":[{\"type\":\"image_generation\",\"model\":\"gpt-image-2\",\"size\":\"auto\",\"quality\":\"auto\",\"output_format\":\"png\"}]}}\n\n"+
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000000,\"tools\":[{\"type\":\"image_generation\",\"model\":\"gpt-image-2\",\"size\":\"auto\",\"quality\":\"auto\",\"output_format\":\"png\"}],\"output\":[{\"id\":\"ig_actual_size\",\"type\":\"image_generation_call\",\"result\":%q}]}}\n\n"+
|
||||
"data: [DONE]\n\n",
|
||||
encoded,
|
||||
)
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"text/event-stream"},
|
||||
"X-Request-Id": []string{"req_img_actual_size"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(upstreamBody)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{httpUpstream: upstream}
|
||||
parsed, err := svc.ParseOpenAIImagesRequest(c, body)
|
||||
require.NoError(t, err)
|
||||
|
||||
account := &Account{
|
||||
ID: 1,
|
||||
Name: "openai-oauth",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "token-123",
|
||||
"chatgpt_account_id": "acct-123",
|
||||
},
|
||||
}
|
||||
result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
return openAIOAuthImageActualSizeTestRun{result: result, recorder: rec, upstream: upstream}
|
||||
}
|
||||
|
||||
func encodeOpenAIImageTestPNG(t *testing.T, width, height int) string {
|
||||
t.Helper()
|
||||
img := image.NewNRGBA(image.Rect(0, 0, width, height))
|
||||
img.SetNRGBA(0, 0, color.NRGBA{R: 0xff, A: 0xff})
|
||||
var buf bytes.Buffer
|
||||
require.NoError(t, png.Encode(&buf, img))
|
||||
return base64.StdEncoding.EncodeToString(buf.Bytes())
|
||||
}
|
||||
|
||||
func encodeOpenAIImageTestJPEG(t *testing.T, width, height int) string {
|
||||
t.Helper()
|
||||
img := image.NewNRGBA(image.Rect(0, 0, width, height))
|
||||
img.SetNRGBA(0, 0, color.NRGBA{G: 0xff, A: 0xff})
|
||||
var buf bytes.Buffer
|
||||
require.NoError(t, jpeg.Encode(&buf, img, nil))
|
||||
return base64.StdEncoding.EncodeToString(buf.Bytes())
|
||||
}
|
||||
|
||||
func encodeOpenAIImageTestWebPVP8X(width, height int) string {
|
||||
header := make([]byte, 30)
|
||||
copy(header[0:4], "RIFF")
|
||||
copy(header[8:12], "WEBP")
|
||||
copy(header[12:16], "VP8X")
|
||||
width--
|
||||
height--
|
||||
header[24], header[25], header[26] = byte(width), byte(width>>8), byte(width>>16)
|
||||
header[27], header[28], header[29] = byte(height), byte(height>>8), byte(height>>16)
|
||||
return base64.StdEncoding.EncodeToString(header)
|
||||
}
|
||||
|
||||
func encodeOpenAIImageTestWebPVP8(width, height int) string {
|
||||
header := make([]byte, 30)
|
||||
copy(header[0:4], "RIFF")
|
||||
copy(header[8:12], "WEBP")
|
||||
copy(header[12:16], "VP8 ")
|
||||
copy(header[23:26], "\x9d\x01\x2a")
|
||||
binary.LittleEndian.PutUint16(header[26:28], uint16(width))
|
||||
binary.LittleEndian.PutUint16(header[28:30], uint16(height))
|
||||
return base64.StdEncoding.EncodeToString(header)
|
||||
}
|
||||
|
||||
func encodeOpenAIImageTestWebPVP8L(width, height int) string {
|
||||
header := make([]byte, 25)
|
||||
copy(header[0:4], "RIFF")
|
||||
copy(header[8:12], "WEBP")
|
||||
copy(header[12:16], "VP8L")
|
||||
header[20] = 0x2f
|
||||
width--
|
||||
height--
|
||||
header[21] = byte(width)
|
||||
header[22] = byte(width>>8)&0x3f | byte(height&0x03)<<6
|
||||
header[23] = byte(height >> 2)
|
||||
header[24] = byte(height>>10) & 0x0f
|
||||
return base64.StdEncoding.EncodeToString(header)
|
||||
}
|
||||
@@ -565,12 +565,14 @@ func collectOpenAIImagesFromResponsesBody(body []byte) ([]openAIResponsesImageRe
|
||||
return nil, 0, nil, openAIResponsesImageResult{}, false, collectErr
|
||||
}
|
||||
if len(finalResults) > 0 {
|
||||
reconcileOpenAIResponsesImageResultSizes(finalResults, &finalMeta)
|
||||
return finalResults, createdAt, usageRaw, finalMeta, true, nil
|
||||
}
|
||||
|
||||
if len(fallbackResults) > 0 {
|
||||
firstMeta := fallbackResults[0]
|
||||
mergeOpenAIResponsesImageMeta(&firstMeta, responseMeta)
|
||||
reconcileOpenAIResponsesImageResultSizes(fallbackResults, &firstMeta)
|
||||
return fallbackResults, createdAt, usageRaw, firstMeta, foundFinal, nil
|
||||
}
|
||||
return nil, createdAt, usageRaw, openAIResponsesImageResult{}, foundFinal, nil
|
||||
@@ -1262,6 +1264,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse(
|
||||
mergeOpenAIResponsesImageMeta(&img, streamMeta)
|
||||
appendOpenAIResponsesImageResultDedup(&finalResults, finalSeen, "", img)
|
||||
}
|
||||
reconcileOpenAIResponsesImageResultSizes(finalResults, nil)
|
||||
if len(finalResults) == 0 {
|
||||
outputErr := fmt.Errorf("upstream did not return image output")
|
||||
// 软失败:response.completed 事件里没有图片。记录上游诊断摘要到 ops,
|
||||
@@ -1324,8 +1327,12 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse(
|
||||
}
|
||||
if len(pendingResults) > 0 {
|
||||
eventName := streamPrefix + ".completed"
|
||||
for _, img := range pendingResults {
|
||||
mergeOpenAIResponsesImageMeta(&img, streamMeta)
|
||||
finalResults := append([]openAIResponsesImageResult(nil), pendingResults...)
|
||||
for i := range finalResults {
|
||||
mergeOpenAIResponsesImageMeta(&finalResults[i], streamMeta)
|
||||
}
|
||||
reconcileOpenAIResponsesImageResultSizes(finalResults, nil)
|
||||
for _, img := range finalResults {
|
||||
key := openAIResponsesImageResultKey("", img)
|
||||
if _, exists := emitted[key]; exists {
|
||||
continue
|
||||
@@ -1335,7 +1342,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse(
|
||||
s.tryWriteOpenAIImagesStreamEvent(c, flusher, &clientDisconnected, &lastDownstreamWriteAt, eventName, payload)
|
||||
}
|
||||
imageCount = len(emitted)
|
||||
imageOutputSizes = openAIResponsesImageResultSizes(pendingResults)
|
||||
imageOutputSizes = openAIResponsesImageResultSizes(finalResults)
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user