Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,753 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestOpenAIGatewayServiceForward_RejectsDisabledImageGenerationIntents(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
body []byte
|
||||
}{
|
||||
{
|
||||
name: "image model",
|
||||
body: []byte(`{"model":"gpt-image-2","input":"draw"}`),
|
||||
},
|
||||
{
|
||||
name: "image tool",
|
||||
body: []byte(`{"model":"gpt-5.4","input":"draw","tools":[{"type":"image_generation"}]}`),
|
||||
},
|
||||
{
|
||||
name: "image tool choice",
|
||||
body: []byte(`{"model":"gpt-5.4","input":"draw","tool_choice":{"type":"image_generation"}}`),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
c, recorder := newOpenAIImageGenerationControlTestContext(false, "unit-test-agent/1.0")
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
|
||||
result, err := svc.Forward(context.Background(), c, account, tt.body)
|
||||
|
||||
require.Error(t, err)
|
||||
require.Nil(t, result)
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code)
|
||||
require.Equal(t, "permission_error", gjson.GetBytes(recorder.Body.Bytes(), "error.type").String())
|
||||
require.Nil(t, upstream.lastReq, "disabled image request must not reach upstream")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForward_DisabledGroupAllowsTextOnlyResponses(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstream := &httpUpstreamRecorder{
|
||||
resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_text","model":"gpt-5.4","usage":{"input_tokens":3,"output_tokens":2}}`)),
|
||||
},
|
||||
}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
c, recorder := newOpenAIImageGenerationControlTestContext(false, "unit-test-agent/1.0")
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
|
||||
result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.4","input":"write code","stream":false}`))
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.Equal(t, 3, result.Usage.InputTokens)
|
||||
require.Equal(t, 2, result.Usage.OutputTokens)
|
||||
require.Equal(t, 0, result.ImageCount)
|
||||
require.NotNil(t, upstream.lastReq)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForward_CodexImageInjectionRespectsGroupCapability(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
allowImages bool
|
||||
bridgeEnabled bool
|
||||
responsesLite bool
|
||||
wantInjected bool
|
||||
}{
|
||||
{name: "disabled group skips injection", allowImages: false, bridgeEnabled: true, wantInjected: false},
|
||||
{name: "enabled group skips injection by default", allowImages: true, bridgeEnabled: false, wantInjected: false},
|
||||
{name: "enabled group injects image tool when bridge enabled", allowImages: true, bridgeEnabled: true, wantInjected: true},
|
||||
{name: "responses lite skips hosted image bridge", allowImages: true, bridgeEnabled: true, responsesLite: true, wantInjected: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{
|
||||
resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_codex","model":"gpt-5.4","usage":{"input_tokens":1,"output_tokens":1}}`)),
|
||||
},
|
||||
}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
svc.cfg.Gateway.CodexImageGenerationBridgeEnabled = tt.bridgeEnabled
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(tt.allowImages, "codex_cli_rs/0.98.0")
|
||||
if tt.responsesLite {
|
||||
c.Request.Header.Set(responsesLiteHeader, "true")
|
||||
}
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
|
||||
result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.4","input":"write code","stream":false}`))
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, upstream.lastReq)
|
||||
hasImageTool := gjson.GetBytes(upstream.lastBody, `tools.#(type=="image_generation")`).Exists()
|
||||
require.Equal(t, tt.wantInjected, hasImageTool)
|
||||
expectedLiteHeader := ""
|
||||
if tt.responsesLite {
|
||||
expectedLiteHeader = "true"
|
||||
}
|
||||
require.Equal(t, expectedLiteHeader, upstream.lastReq.Header.Get(responsesLiteHeader))
|
||||
instructions := gjson.GetBytes(upstream.lastBody, "instructions").String()
|
||||
require.Equal(t, tt.wantInjected, strings.Contains(instructions, "image_generation"))
|
||||
toolChoice := gjson.GetBytes(upstream.lastBody, "tool_choice")
|
||||
require.Equal(t, tt.wantInjected, toolChoice.Exists())
|
||||
if tt.wantInjected {
|
||||
require.Equal(t, "auto", toolChoice.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIBuildUpstreamRequestOpenAIPassthroughForwardsResponsesLiteHeader(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(true, "codex_cli_rs/0.98.0")
|
||||
c.Request.Header.Set(responsesLiteHeader, "true")
|
||||
|
||||
svc := newOpenAIImageGenerationControlTestService(&httpUpstreamRecorder{})
|
||||
req, err := svc.buildUpstreamRequestOpenAIPassthrough(
|
||||
c.Request.Context(),
|
||||
c,
|
||||
newOpenAIImageGenerationControlTestAccount(),
|
||||
[]byte(`{"model":"gpt-5.4","input":"write code"}`),
|
||||
"test-token",
|
||||
)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "true", req.Header.Get(responsesLiteHeader))
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForward_ExplicitImageToolWorksWithBridgeDisabled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstream := &httpUpstreamRecorder{
|
||||
resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_explicit_image","model":"gpt-5.4","usage":{"input_tokens":2,"output_tokens":1}}`)),
|
||||
},
|
||||
}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(true, "codex_cli_rs/0.98.0")
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
body := []byte(`{"model":"gpt-5.4","input":"draw","stream":false,"tools":[{"type":"image_generation","format":"jpeg"}]}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), c, account, body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, upstream.lastReq)
|
||||
require.True(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="image_generation")`).Exists())
|
||||
require.Equal(t, "jpeg", gjson.GetBytes(upstream.lastBody, `tools.#(type=="image_generation").output_format`).String())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="image_generation").format`).Exists())
|
||||
instructions := gjson.GetBytes(upstream.lastBody, "instructions").String()
|
||||
require.NotContains(t, instructions, "image_generation")
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForward_AccountPolicyStripsExplicitImageTool(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstream := &httpUpstreamRecorder{
|
||||
resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_stripped_image","model":"gpt-5.4","usage":{"input_tokens":2,"output_tokens":1}}`)),
|
||||
},
|
||||
}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(true, "codex_cli_rs/0.98.0")
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
account.Extra = map[string]any{
|
||||
featureKeyCodexImageGenerationExplicitToolPolicy: codexImageGenerationExplicitToolPolicyStrip,
|
||||
}
|
||||
body := []byte(`{
|
||||
"model":"gpt-5.4",
|
||||
"input":"draw",
|
||||
"stream":false,
|
||||
"tools":[
|
||||
{"type":"function","name":"shell","parameters":{"type":"object"}},
|
||||
{"type":"image_generation","format":"jpeg"}
|
||||
],
|
||||
"tool_choice":{"type":"image_generation"}
|
||||
}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), c, account, body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, upstream.lastReq)
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="image_generation")`).Exists())
|
||||
require.True(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="function")`).Exists())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists())
|
||||
instructions := gjson.GetBytes(upstream.lastBody, "instructions").String()
|
||||
require.NotContains(t, instructions, "image_generation")
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForward_AccountPolicyStripsImageNamespaceTools(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
passthrough bool
|
||||
}{
|
||||
{name: "managed forwarding"},
|
||||
{name: "passthrough forwarding", passthrough: true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{
|
||||
resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_stripped_namespace","model":"gpt-5.5","usage":{"input_tokens":2,"output_tokens":1}}`)),
|
||||
},
|
||||
}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(false, "codex_cli_rs/0.144.1")
|
||||
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
account.Extra = map[string]any{
|
||||
featureKeyCodexImageGenerationExplicitToolPolicy: codexImageGenerationExplicitToolPolicyStrip,
|
||||
"openai_passthrough": tt.passthrough,
|
||||
}
|
||||
body := []byte(`{
|
||||
"model":"gpt-5.5",
|
||||
"stream":false,
|
||||
"tools":[
|
||||
{"type":"function","name":"shell","parameters":{"type":"object"}},
|
||||
{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]},
|
||||
{"type":"namespace","name":"code_tools","tools":[{"type":"function","name":"run"}]}
|
||||
],
|
||||
"input":[
|
||||
{"type":"message","role":"user","content":[{"type":"input_text","text":"write code"}]},
|
||||
{"type":"additional_tools","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]}
|
||||
],
|
||||
"tool_choice":"auto"
|
||||
}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), c, account, body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, upstream.lastReq)
|
||||
var forwarded map[string]any
|
||||
require.NoError(t, json.Unmarshal(upstream.lastBody, &forwarded))
|
||||
require.False(t, hasOpenAIImageGenerationTool(forwarded))
|
||||
require.Equal(t, "auto", forwarded["tool_choice"])
|
||||
require.True(t, gjson.GetBytes(upstream.lastBody, `tools.#(name=="shell")`).Exists())
|
||||
require.True(t, gjson.GetBytes(upstream.lastBody, `tools.#(name=="code_tools")`).Exists())
|
||||
require.Equal(t, "write code", gjson.GetBytes(upstream.lastBody, "input.0.content.0.text").String())
|
||||
cached, known := getOpenAIImageIntentHint(c)
|
||||
require.True(t, known)
|
||||
require.True(t, cached)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForward_ChannelBridgeOverrideEnablesCodexInjection(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstream := &httpUpstreamRecorder{
|
||||
resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_channel_bridge","model":"gpt-5.4","usage":{"input_tokens":1,"output_tokens":1}}`)),
|
||||
},
|
||||
}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
groupID := int64(4242)
|
||||
svc.channelService = newOpenAIImageGenerationControlChannelService(groupID, &Channel{
|
||||
ID: 9001,
|
||||
Status: StatusActive,
|
||||
FeaturesConfig: map[string]any{
|
||||
featureKeyCodexImageGenerationBridge: map[string]any{PlatformOpenAI: true},
|
||||
},
|
||||
})
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(true, "codex_cli_rs/0.98.0")
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
|
||||
result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.4","input":"write code","stream":false}`))
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, upstream.lastReq)
|
||||
require.True(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="image_generation")`).Exists())
|
||||
require.Equal(t, "auto", gjson.GetBytes(upstream.lastBody, "tool_choice").String())
|
||||
instructions := gjson.GetBytes(upstream.lastBody, "instructions").String()
|
||||
require.Contains(t, instructions, "image_generation")
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForward_CodexBridgeDoesNotInjectHostedToolAlongsideImageGenNamespace(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstream := &httpUpstreamRecorder{
|
||||
resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_namespace_image","model":"gpt-5.5","usage":{"input_tokens":1,"output_tokens":1}}`)),
|
||||
},
|
||||
}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
svc.cfg.Gateway.CodexImageGenerationBridgeEnabled = true
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(true, "codex_cli_rs/0.144.1")
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
body := []byte(`{
|
||||
"model":"gpt-5.5",
|
||||
"stream":false,
|
||||
"tools":[
|
||||
{"type":"function","name":"shell","parameters":{"type":"object"}},
|
||||
{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}
|
||||
],
|
||||
"input":[
|
||||
{"type":"message","role":"user","content":[{"type":"input_text","text":"draw a cat"}]},
|
||||
{"type":"additional_tools","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]}
|
||||
],
|
||||
"tool_choice":"auto"
|
||||
}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), c, account, body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, upstream.lastReq)
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="image_generation")`).Exists())
|
||||
require.Equal(t, "namespace", gjson.GetBytes(upstream.lastBody, `tools.#(name=="image_gen").type`).String())
|
||||
require.Equal(t, "namespace", gjson.GetBytes(upstream.lastBody, `input.#(type=="additional_tools").tools.#(name=="image_gen").type`).String())
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForward_CodexBridgePreservesImageGenFunction(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
tool string
|
||||
}{
|
||||
{
|
||||
name: "flat function",
|
||||
tool: `{"type":"function","name":"image_gen.imagegen","parameters":{"type":"object"}}`,
|
||||
},
|
||||
{
|
||||
name: "nested function",
|
||||
tool: `{"type":"function","function":{"name":"image_gen.imagegen","parameters":{"type":"object"}}}`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{
|
||||
resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_function_image","model":"gpt-5.5","usage":{"input_tokens":1,"output_tokens":1}}`)),
|
||||
},
|
||||
}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
svc.cfg.Gateway.CodexImageGenerationBridgeEnabled = true
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(true, "codex_cli_rs/0.144.1")
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
body := []byte(`{"model":"gpt-5.5","input":"draw a cat","stream":false,"tools":[` + tt.tool + `]}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), c, account, body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, upstream.lastReq)
|
||||
|
||||
var forwarded map[string]any
|
||||
require.NoError(t, json.Unmarshal(upstream.lastBody, &forwarded))
|
||||
require.True(t, hasCodexImageGenerationFunctionTool(forwarded))
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="image_generation")`).Exists())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists())
|
||||
require.NotContains(t, gjson.GetBytes(upstream.lastBody, "instructions").String(), codexImageGenerationBridgeMarker)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForward_CodexBridgePreservesExistingToolChoice(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstream := &httpUpstreamRecorder{
|
||||
resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_codex_tool_choice","model":"gpt-5.4","usage":{"input_tokens":1,"output_tokens":1}}`)),
|
||||
},
|
||||
}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
svc.cfg.Gateway.CodexImageGenerationBridgeEnabled = true
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(true, "codex_cli_rs/0.98.0")
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
|
||||
result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.4","input":"draw","stream":false,"tools":[{"type":"image_generation"}],"tool_choice":{"type":"image_generation"}}`))
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, upstream.lastReq)
|
||||
require.Equal(t, "image_generation", gjson.GetBytes(upstream.lastBody, "tool_choice.type").String())
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForward_CodexBridgeSkipsCompactRequests(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstream := &httpUpstreamRecorder{
|
||||
resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_codex_compact","model":"gpt-5.4","usage":{"input_tokens":1,"output_tokens":1}}`)),
|
||||
},
|
||||
}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
svc.cfg.Gateway.CodexImageGenerationBridgeEnabled = true
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(true, "codex_cli_rs/0.98.0")
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses/compact", nil)
|
||||
c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
|
||||
// /responses/compact 上游不接受 tool_choice,bridge 注入必须整体豁免 compact 请求。
|
||||
result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.4","input":"summarize the conversation","stream":false}`))
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, upstream.lastReq)
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="image_generation")`).Exists())
|
||||
instructions := gjson.GetBytes(upstream.lastBody, "instructions").String()
|
||||
require.NotContains(t, instructions, "image_generation")
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_CodexImageGenerationBridgeOverridePrecedence(t *testing.T) {
|
||||
groupID := int64(4242)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
global bool
|
||||
channel *Channel
|
||||
account *Account
|
||||
want bool
|
||||
}{
|
||||
{
|
||||
name: "global default enables bridge",
|
||||
global: true,
|
||||
account: &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
},
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "channel true overrides disabled global",
|
||||
global: false,
|
||||
channel: &Channel{ID: 1, Status: StatusActive, FeaturesConfig: map[string]any{
|
||||
featureKeyCodexImageGenerationBridge: map[string]any{PlatformOpenAI: true},
|
||||
}},
|
||||
account: &Account{Platform: PlatformOpenAI},
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "channel false overrides enabled global",
|
||||
global: true,
|
||||
channel: &Channel{ID: 1, Status: StatusActive, FeaturesConfig: map[string]any{
|
||||
featureKeyCodexImageGenerationBridge: map[string]any{PlatformOpenAI: false},
|
||||
}},
|
||||
account: &Account{Platform: PlatformOpenAI},
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "account false overrides channel and global true",
|
||||
global: true,
|
||||
channel: &Channel{ID: 1, Status: StatusActive, FeaturesConfig: map[string]any{
|
||||
featureKeyCodexImageGenerationBridge: map[string]any{PlatformOpenAI: true},
|
||||
}},
|
||||
account: &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Extra: map[string]any{featureKeyCodexImageGenerationBridge: false},
|
||||
},
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "nested account true overrides channel false",
|
||||
global: false,
|
||||
channel: &Channel{ID: 1, Status: StatusActive, FeaturesConfig: map[string]any{
|
||||
featureKeyCodexImageGenerationBridge: map[string]any{PlatformOpenAI: false},
|
||||
}},
|
||||
account: &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Extra: map[string]any{
|
||||
PlatformOpenAI: map[string]any{"codex_image_generation_bridge_enabled": true},
|
||||
},
|
||||
},
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "non openai account extra is ignored",
|
||||
global: false,
|
||||
account: &Account{
|
||||
Platform: PlatformAnthropic,
|
||||
Extra: map[string]any{featureKeyCodexImageGenerationBridge: true},
|
||||
},
|
||||
want: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
svc := newOpenAIImageGenerationControlTestService(&httpUpstreamRecorder{})
|
||||
svc.cfg.Gateway.CodexImageGenerationBridgeEnabled = tt.global
|
||||
if tt.channel != nil {
|
||||
svc.channelService = newOpenAIImageGenerationControlChannelService(groupID, tt.channel)
|
||||
}
|
||||
apiKey := &APIKey{GroupID: &groupID}
|
||||
|
||||
got := svc.isCodexImageGenerationBridgeEnabled(context.Background(), tt.account, apiKey)
|
||||
|
||||
require.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceHandleResponsesImageOutputs_NonStreaming(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
svc := newOpenAIImageGenerationControlTestService(&httpUpstreamRecorder{})
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(true, "unit-test-agent/1.0")
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{
|
||||
"id":"resp_image_json",
|
||||
"model":"gpt-5.4",
|
||||
"output":[{"id":"ig_json_1","type":"image_generation_call","result":"final-image"}],
|
||||
"usage":{"input_tokens":7,"output_tokens":3,"output_tokens_details":{"image_tokens":2}}
|
||||
}`)),
|
||||
}
|
||||
|
||||
result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, &Account{ID: 1, Type: AccountTypeAPIKey}, "gpt-5.4", "gpt-5.4")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 1, result.imageCount)
|
||||
require.NotNil(t, result.usage)
|
||||
require.Equal(t, 7, result.usage.InputTokens)
|
||||
require.Equal(t, 3, result.usage.OutputTokens)
|
||||
require.Equal(t, 2, result.usage.ImageOutputTokens)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceHandleResponsesImageOutputs_Streaming(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
svc := newOpenAIImageGenerationControlTestService(&httpUpstreamRecorder{})
|
||||
c, recorder := newOpenAIImageGenerationControlTestContext(true, "unit-test-agent/1.0")
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
"data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"ig_stream_1\",\"type\":\"image_generation_call\",\"status\":\"generating\",\"result\":\"final-image\"}}\n\n" +
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_image_stream\",\"model\":\"gpt-5.5\",\"output\":[{\"id\":\"ig_stream_1\",\"type\":\"image_generation_call\",\"status\":\"generating\",\"result\":\"final-image\"}],\"usage\":{\"input_tokens\":11,\"output_tokens\":5,\"output_tokens_details\":{\"image_tokens\":4}}}}\n\n",
|
||||
)),
|
||||
}
|
||||
|
||||
result, err := svc.handleStreamingResponse(context.Background(), resp, c, &Account{ID: 1}, time.Now(), "gpt-5.5", "gpt-5.5")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 1, result.imageCount)
|
||||
require.NotNil(t, result.usage)
|
||||
require.Equal(t, 11, result.usage.InputTokens)
|
||||
require.Equal(t, 5, result.usage.OutputTokens)
|
||||
require.Equal(t, 4, result.usage.ImageOutputTokens)
|
||||
require.NotContains(t, recorder.Body.String(), `"status":"generating"`)
|
||||
require.Equal(t, 2, strings.Count(recorder.Body.String(), `"status":"completed"`))
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceHandleResponsesImageOutputs_StreamingPassthrough(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
svc := newOpenAIImageGenerationControlTestService(&httpUpstreamRecorder{})
|
||||
c, recorder := newOpenAIImageGenerationControlTestContext(true, "unit-test-agent/1.0")
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
"data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"ig_stream_1\",\"type\":\"image_generation_call\",\"status\":\"in_progress\",\"result\":\"final-image\"}}\n\n" +
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_image_stream\",\"model\":\"gpt-5.5\",\"output\":[{\"id\":\"ig_stream_1\",\"type\":\"image_generation_call\",\"status\":\"in_progress\",\"result\":\"final-image\"}],\"usage\":{\"input_tokens\":11,\"output_tokens\":5,\"output_tokens_details\":{\"image_tokens\":4}}}}\n\n",
|
||||
)),
|
||||
}
|
||||
|
||||
result, err := svc.handleStreamingResponsePassthrough(context.Background(), resp, c, &Account{ID: 1}, time.Now(), "gpt-5.5", "gpt-5.5")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotContains(t, recorder.Body.String(), `"status":"in_progress"`)
|
||||
require.Equal(t, 2, strings.Count(recorder.Body.String(), `"status":"completed"`))
|
||||
}
|
||||
|
||||
func TestNormalizeCompletedImageGenerationStatus(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
want string
|
||||
wantChanged bool
|
||||
}{
|
||||
{
|
||||
name: "output item done with result",
|
||||
input: `{"type":"response.output_item.done","item":{"type":"image_generation_call","status":"generating","result":"image-data"}}`,
|
||||
want: `{"type":"response.output_item.done","item":{"type":"image_generation_call","status":"completed","result":"image-data"}}`,
|
||||
wantChanged: true,
|
||||
},
|
||||
{
|
||||
name: "terminal response only changes completed image result",
|
||||
input: `{"type":"response.completed","response":{"output":[{"type":"image_generation_call","status":"in_progress","result":"image-data"},{"type":"image_generation_call","status":"failed","result":"partial-data"}]}}`,
|
||||
want: `{"type":"response.completed","response":{"output":[{"type":"image_generation_call","status":"completed","result":"image-data"},{"type":"image_generation_call","status":"failed","result":"partial-data"}]}}`,
|
||||
wantChanged: true,
|
||||
},
|
||||
{
|
||||
name: "done item without result",
|
||||
input: `{"type":"response.output_item.done","item":{"type":"image_generation_call","status":"generating"}}`,
|
||||
want: `{"type":"response.output_item.done","item":{"type":"image_generation_call","status":"generating"}}`,
|
||||
wantChanged: false,
|
||||
},
|
||||
{
|
||||
name: "non-final image event",
|
||||
input: `{"type":"response.output_item.added","item":{"type":"image_generation_call","status":"generating","result":"image-data"}}`,
|
||||
want: `{"type":"response.output_item.added","item":{"type":"image_generation_call","status":"generating","result":"image-data"}}`,
|
||||
wantChanged: false,
|
||||
},
|
||||
{
|
||||
name: "done preserves base64 result",
|
||||
input: `{"type":"response.done","response":{"output":[{"type":"image_generation_call","status":"generating","result":"iVBORw0KGgoAAAANSUhEUg/+=="}]}}`,
|
||||
want: `{"type":"response.done","response":{"output":[{"type":"image_generation_call","status":"completed","result":"iVBORw0KGgoAAAANSUhEUg/+=="}]}}`,
|
||||
wantChanged: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, changed := normalizeCompletedImageGenerationStatus([]byte(tt.input))
|
||||
|
||||
require.Equal(t, tt.wantChanged, changed)
|
||||
require.JSONEq(t, tt.want, string(got))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleStreamingResponse_CyberPolicyCapturesRealUpstreamTokens 锁定流式
|
||||
// /v1/responses 命中 cyber_policy 的计费正确性:response.failed 自带的真实 usage
|
||||
// 必须在打 cyber 标记前被解析进 mark;否则计费走 mark.UpstreamInTok 会按 0 token
|
||||
// 漏记真实用量(该路径返回错误,handler 仅经 RecordCyberPolicyUsageLog 计费)。
|
||||
func TestHandleStreamingResponse_CyberPolicyCapturesRealUpstreamTokens(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
svc := newOpenAIImageGenerationControlTestService(&httpUpstreamRecorder{})
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(false, "unit-test-agent/1.0")
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_cyber\"}}\n\n" +
|
||||
"data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_cyber\",\"error\":{\"code\":\"cyber_policy\",\"message\":\"blocked by network policy\"},\"usage\":{\"input_tokens\":1234,\"output_tokens\":7}}}\n\n",
|
||||
)),
|
||||
}
|
||||
|
||||
_, err := svc.handleStreamingResponse(context.Background(), resp, c, &Account{ID: 1}, time.Now(), "gpt-5.5", "gpt-5.5")
|
||||
require.Error(t, err, "cyber 命中的流式响应应返回错误(sawFailedEvent)")
|
||||
|
||||
mark := GetOpsCyberPolicy(c)
|
||||
require.NotNil(t, mark, "必须打上 cyber 标记")
|
||||
require.Equal(t, "cyber_policy", mark.Code)
|
||||
require.Equal(t, 1234, mark.UpstreamInTok, "必须捕获 response.failed 自带真实 input token,而非解析前的 0")
|
||||
require.Equal(t, 7, mark.UpstreamOutTok)
|
||||
}
|
||||
|
||||
func newOpenAIImageGenerationControlTestService(upstream *httpUpstreamRecorder) *OpenAIGatewayService {
|
||||
cfg := &config.Config{}
|
||||
return &OpenAIGatewayService{
|
||||
cfg: cfg,
|
||||
httpUpstream: upstream,
|
||||
cache: &stubGatewayCache{},
|
||||
openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
|
||||
toolCorrector: NewCodexToolCorrector(),
|
||||
}
|
||||
}
|
||||
|
||||
func newOpenAIImageGenerationControlChannelService(groupID int64, ch *Channel) *ChannelService {
|
||||
svc := &ChannelService{}
|
||||
cache := newEmptyChannelCache()
|
||||
if ch != nil {
|
||||
cache.channelByGroupID[groupID] = ch
|
||||
cache.byID[ch.ID] = ch
|
||||
}
|
||||
cache.loadedAt = time.Now()
|
||||
svc.cache.Store(cache)
|
||||
return svc
|
||||
}
|
||||
|
||||
func newOpenAIImageGenerationControlTestContext(allowImages bool, userAgent string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
|
||||
c.Request.Header.Set("User-Agent", userAgent)
|
||||
groupID := int64(4242)
|
||||
c.Set("api_key", &APIKey{
|
||||
ID: 2424,
|
||||
GroupID: &groupID,
|
||||
Group: &Group{
|
||||
ID: groupID,
|
||||
AllowImageGeneration: allowImages,
|
||||
RateMultiplier: 1,
|
||||
ImageRateMultiplier: 1,
|
||||
},
|
||||
})
|
||||
return c, recorder
|
||||
}
|
||||
|
||||
func newOpenAIImageGenerationControlTestAccount() *Account {
|
||||
return &Account{
|
||||
ID: 5151,
|
||||
Name: "openai-image-controls",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeAPIKey,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "sk-test",
|
||||
},
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user