| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113 |
- package handler
- import (
- "net/http"
- "net/http/httptest"
- "strings"
- "testing"
- "github.com/Wei-Shaw/sub2api/internal/config"
- middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
- "github.com/Wei-Shaw/sub2api/internal/service"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/require"
- )
- func TestOpenAIGatewayHandlerResponses_GrokPassiveImageToolDeclarationBypassesPermissionGate(t *testing.T) {
- body := `{"model":"grok-4.5","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}],"tool_choice":"auto","input":"write code"}`
- rec := runOpenAIResponsesImagePermissionGateTest(t, service.PlatformGrok, body)
- require.NotEqual(t, http.StatusForbidden, rec.Code)
- require.NotContains(t, rec.Body.String(), service.ImageGenerationPermissionMessage())
- }
- func TestOpenAIGatewayHandlerResponses_GrokResponsesLiteImageToolDeclarationBypassesPermissionGate(t *testing.T) {
- body := `{"model":"grok-4.5","tool_choice":"auto","input":[{"type":"additional_tools","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]},{"type":"message","role":"user","content":"write code"}]}`
- rec := runOpenAIResponsesImagePermissionGateTest(t, service.PlatformGrok, body)
- require.NotEqual(t, http.StatusForbidden, rec.Code)
- require.NotContains(t, rec.Body.String(), service.ImageGenerationPermissionMessage())
- }
- func TestOpenAIGatewayHandlerResponses_ImagePermissionHardSignalsStillRejected(t *testing.T) {
- tests := []struct {
- name string
- platform string
- body string
- }{
- {
- name: "Grok native image_generation declaration",
- platform: service.PlatformGrok,
- body: `{"model":"grok-4.5","tools":[{"type":"image_generation"}],"input":"draw"}`,
- },
- {
- name: "Grok explicit image_gen tool choice",
- platform: service.PlatformGrok,
- body: `{"model":"grok-4.5","tools":[{"type":"namespace","name":"image_gen"}],"tool_choice":{"type":"namespace","name":"image_gen"},"input":"draw"}`,
- },
- {
- name: "OpenAI native image_generation tool",
- platform: service.PlatformOpenAI,
- body: `{"model":"gpt-5.5","tools":[{"type":"image_generation","model":"gpt-image-2"}],"input":"draw a cat"}`,
- },
- {
- name: "OpenAI image model",
- platform: service.PlatformOpenAI,
- body: `{"model":"gpt-image-2","input":"draw a cat"}`,
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- rec := runOpenAIResponsesImagePermissionGateTest(t, tt.platform, tt.body)
- require.Equal(t, http.StatusForbidden, rec.Code)
- require.Contains(t, rec.Body.String(), service.ImageGenerationPermissionMessage())
- })
- }
- }
- func TestOpenAIGatewayHandlerResponses_PassiveNamespaceDoesNotTrigger403(t *testing.T) {
- passiveNamespace := `{"model":"gpt-5.5","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}],"tool_choice":"auto","input":"write code"}`
- rec := runOpenAIResponsesImagePermissionGateTest(t, service.PlatformOpenAI, passiveNamespace)
- require.NotEqual(t, http.StatusForbidden, rec.Code,
- "passive image_gen namespace with tool_choice=auto should not trigger 403 (#4447)")
- }
- func runOpenAIResponsesImagePermissionGateTest(t *testing.T, platform string, body string) *httptest.ResponseRecorder {
- t.Helper()
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(body))
- c.Request.Header.Set("Content-Type", "application/json")
- groupID := int64(6301)
- userID := int64(6302)
- c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
- ID: 6303,
- GroupID: &groupID,
- Group: &service.Group{
- ID: groupID,
- Platform: platform,
- AllowImageGeneration: false,
- },
- User: &service.User{ID: userID, Status: service.StatusActive},
- })
- c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: userID, Concurrency: 1})
- h := &OpenAIGatewayHandler{
- gatewayService: &service.OpenAIGatewayService{},
- billingCacheService: service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, &config.Config{RunMode: config.RunModeSimple}, nil),
- apiKeyService: &service.APIKeyService{},
- concurrencyHelper: &ConcurrencyHelper{concurrencyService: service.NewConcurrencyService(
- &helperConcurrencyCacheStub{userSeq: []bool{true}},
- )},
- cfg: &config.Config{},
- imageLimiter: &imageConcurrencyLimiter{},
- }
- h.Responses(c)
- return rec
- }
|