feat: add tri-state support for disable-image-generation configuration

- Introduced `DisableImageGenerationMode` with support for `false`, `true`, and `chat` values.
- Updated payload handling to preserve `image_generation` on images endpoints when `chat` mode is enabled.
- Modified OpenAI image handlers (`ImagesGenerations`, `ImagesEdits`) to respect tri-state logic.
- Added unit tests for `DisableImageGenerationMode` behavior and endpoint-specific handling.
- Enhanced configuration diff logging to support `DisableImageGenerationMode`.
This commit is contained in:
Luis Pater
2026-04-30 11:59:50 +08:00
parent 46018417ad
commit f56a19e5b8
24 changed files with 398 additions and 54 deletions
@@ -14,6 +14,7 @@ import (
"time"
"github.com/gin-gonic/gin"
internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
log "github.com/sirupsen/logrus"
@@ -198,7 +199,7 @@ func parseBoolField(raw string, fallback bool) bool {
}
func (h *OpenAIAPIHandler) ImagesGenerations(c *gin.Context) {
if h != nil && h.BaseAPIHandler != nil && h.BaseAPIHandler.Cfg != nil && h.BaseAPIHandler.Cfg.DisableImageGeneration {
if h != nil && h.BaseAPIHandler != nil && h.BaseAPIHandler.Cfg != nil && h.BaseAPIHandler.Cfg.DisableImageGeneration == internalconfig.DisableImageGenerationAll {
c.AbortWithStatus(http.StatusNotFound)
return
}
@@ -286,7 +287,7 @@ func (h *OpenAIAPIHandler) ImagesGenerations(c *gin.Context) {
}
func (h *OpenAIAPIHandler) ImagesEdits(c *gin.Context) {
if h != nil && h.BaseAPIHandler != nil && h.BaseAPIHandler.Cfg != nil && h.BaseAPIHandler.Cfg.DisableImageGeneration {
if h != nil && h.BaseAPIHandler != nil && h.BaseAPIHandler.Cfg != nil && h.BaseAPIHandler.Cfg.DisableImageGeneration == internalconfig.DisableImageGenerationAll {
c.AbortWithStatus(http.StatusNotFound)
return
}
@@ -10,6 +10,7 @@ import (
"testing"
"github.com/gin-gonic/gin"
internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
"github.com/tidwall/gjson"
@@ -97,7 +98,7 @@ func TestImagesEditsMultipartRejectsUnsupportedModel(t *testing.T) {
}
func TestImagesGenerations_DisableImageGeneration_Returns404(t *testing.T) {
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{DisableImageGeneration: true}, nil)
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{DisableImageGeneration: internalconfig.DisableImageGenerationAll}, nil)
handler := NewOpenAIAPIHandler(base)
body := strings.NewReader(`{"prompt":"draw a square"}`)
@@ -109,7 +110,7 @@ func TestImagesGenerations_DisableImageGeneration_Returns404(t *testing.T) {
}
func TestImagesEdits_DisableImageGeneration_Returns404(t *testing.T) {
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{DisableImageGeneration: true}, nil)
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{DisableImageGeneration: internalconfig.DisableImageGenerationAll}, nil)
handler := NewOpenAIAPIHandler(base)
body := strings.NewReader(`{"prompt":"edit this","images":[{"image_url":"data:image/png;base64,AA=="}]}`)
@@ -119,3 +120,27 @@ func TestImagesEdits_DisableImageGeneration_Returns404(t *testing.T) {
t.Fatalf("status = %d, want %d: %s", resp.Code, http.StatusNotFound, resp.Body.String())
}
}
func TestImagesGenerations_DisableImageGenerationChat_DoesNotReturn404(t *testing.T) {
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{DisableImageGeneration: internalconfig.DisableImageGenerationChat}, nil)
handler := NewOpenAIAPIHandler(base)
body := strings.NewReader(`{"model":"gpt-5.4-mini","prompt":"draw a square"}`)
resp := performImagesEndpointRequest(t, imagesGenerationsPath, "application/json", body, handler.ImagesGenerations)
if resp.Code != http.StatusBadRequest {
t.Fatalf("status = %d, want %d: %s", resp.Code, http.StatusBadRequest, resp.Body.String())
}
}
func TestImagesEdits_DisableImageGenerationChat_DoesNotReturn404(t *testing.T) {
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{DisableImageGeneration: internalconfig.DisableImageGenerationChat}, nil)
handler := NewOpenAIAPIHandler(base)
body := strings.NewReader(`{"model":"gpt-5.4-mini","prompt":"edit this","images":[{"image_url":"data:image/png;base64,AA=="}]}`)
resp := performImagesEndpointRequest(t, imagesEditsPath, "application/json", body, handler.ImagesEdits)
if resp.Code != http.StatusBadRequest {
t.Fatalf("status = %d, want %d: %s", resp.Code, http.StatusBadRequest, resp.Body.String())
}
}