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", }, } }