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`.
2026-04-30 11:59:50 +08:00
parent 46018417ad
commit f56a19e5b8
24 changed files with 398 additions and 54 deletions
@@ -610,7 +610,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
 	cfg.ErrorLogsMaxFiles = 10
 	cfg.UsageStatisticsEnabled = false
 	cfg.DisableCooling = false
-	cfg.DisableImageGeneration = false
+	cfg.DisableImageGeneration = DisableImageGenerationOff
 	cfg.Pprof.Enable = false
 	cfg.Pprof.Addr = DefaultPprofAddr
 	cfg.AmpCode.RestrictManagementToLocalhost = false // Default to false: API key auth is sufficient
@@ -0,0 +1,136 @@
+package config
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"strings"
+
+	"gopkg.in/yaml.v3"
+)
+
+// DisableImageGenerationMode is a tri-state config value for disable-image-generation.
+//
+// It supports:
+//   - false: enabled
+//   - true: disabled everywhere (including /v1/images/* endpoints)
+//   - "chat": disabled for all non-images endpoints, but enabled for /v1/images/generations and /v1/images/edits
+type DisableImageGenerationMode int
+
+const (
+	DisableImageGenerationOff DisableImageGenerationMode = iota
+	DisableImageGenerationAll
+	DisableImageGenerationChat
+)
+
+func (m DisableImageGenerationMode) String() string {
+	switch m {
+	case DisableImageGenerationOff:
+		return "false"
+	case DisableImageGenerationAll:
+		return "true"
+	case DisableImageGenerationChat:
+		return "chat"
+	default:
+		return "false"
+	}
+}
+
+func (m DisableImageGenerationMode) MarshalYAML() (any, error) {
+	switch m {
+	case DisableImageGenerationAll:
+		return true, nil
+	case DisableImageGenerationChat:
+		return "chat", nil
+	default:
+		return false, nil
+	}
+}
+
+func (m *DisableImageGenerationMode) UnmarshalYAML(value *yaml.Node) error {
+	mode, err := parseDisableImageGenerationNode(value)
+	if err != nil {
+		return err
+	}
+	*m = mode
+	return nil
+}
+
+func (m DisableImageGenerationMode) MarshalJSON() ([]byte, error) {
+	switch m {
+	case DisableImageGenerationAll:
+		return []byte("true"), nil
+	case DisableImageGenerationChat:
+		return json.Marshal("chat")
+	default:
+		return []byte("false"), nil
+	}
+}
+
+func (m *DisableImageGenerationMode) UnmarshalJSON(data []byte) error {
+	mode, err := parseDisableImageGenerationJSON(data)
+	if err != nil {
+		return err
+	}
+	*m = mode
+	return nil
+}
+
+func parseDisableImageGenerationNode(value *yaml.Node) (DisableImageGenerationMode, error) {
+	if value == nil {
+		return DisableImageGenerationOff, nil
+	}
+
+	// First try a typed bool decode (covers unquoted true/false and YAML 1.1 bools).
+	var b bool
+	if err := value.Decode(&b); err == nil && value.Kind == yaml.ScalarNode && value.ShortTag() == "!!bool" {
+		if b {
+			return DisableImageGenerationAll, nil
+		}
+		return DisableImageGenerationOff, nil
+	}
+
+	// Fall back to string decoding (covers quoted "true"/"false" and "chat").
+	var s string
+	if err := value.Decode(&s); err != nil {
+		return DisableImageGenerationOff, fmt.Errorf("invalid disable-image-generation value")
+	}
+	return parseDisableImageGenerationString(s)
+}
+
+func parseDisableImageGenerationJSON(data []byte) (DisableImageGenerationMode, error) {
+	trimmed := bytes.TrimSpace(data)
+	if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) {
+		return DisableImageGenerationOff, nil
+	}
+
+	// bool
+	var b bool
+	if err := json.Unmarshal(trimmed, &b); err == nil {
+		if b {
+			return DisableImageGenerationAll, nil
+		}
+		return DisableImageGenerationOff, nil
+	}
+
+	// string
+	var s string
+	if err := json.Unmarshal(trimmed, &s); err != nil {
+		return DisableImageGenerationOff, fmt.Errorf("invalid disable-image-generation value")
+	}
+	return parseDisableImageGenerationString(s)
+}
+
+func parseDisableImageGenerationString(s string) (DisableImageGenerationMode, error) {
+	s = strings.TrimSpace(strings.ToLower(s))
+	switch s {
+	case "", "false", "0", "off", "no":
+		return DisableImageGenerationOff, nil
+	case "true", "1", "on", "yes":
+		return DisableImageGenerationAll, nil
+	case "chat":
+		return DisableImageGenerationChat, nil
+	default:
+		return DisableImageGenerationOff, fmt.Errorf("invalid disable-image-generation value %q (allowed: true, false, chat)", s)
+	}
+}
@@ -0,0 +1,76 @@
+package config
+
+import (
+	"encoding/json"
+	"testing"
+
+	"gopkg.in/yaml.v3"
+)
+
+func TestDisableImageGenerationMode_UnmarshalYAML(t *testing.T) {
+	type wrapper struct {
+		V DisableImageGenerationMode `yaml:"disable-image-generation"`
+	}
+
+	{
+		var w wrapper
+		if err := yaml.Unmarshal([]byte("disable-image-generation: false\n"), &w); err != nil {
+			t.Fatalf("unmarshal false: %v", err)
+		}
+		if w.V != DisableImageGenerationOff {
+			t.Fatalf("false => %v, want %v", w.V, DisableImageGenerationOff)
+		}
+	}
+
+	{
+		var w wrapper
+		if err := yaml.Unmarshal([]byte("disable-image-generation: true\n"), &w); err != nil {
+			t.Fatalf("unmarshal true: %v", err)
+		}
+		if w.V != DisableImageGenerationAll {
+			t.Fatalf("true => %v, want %v", w.V, DisableImageGenerationAll)
+		}
+	}
+
+	{
+		var w wrapper
+		if err := yaml.Unmarshal([]byte("disable-image-generation: chat\n"), &w); err != nil {
+			t.Fatalf("unmarshal chat: %v", err)
+		}
+		if w.V != DisableImageGenerationChat {
+			t.Fatalf("chat => %v, want %v", w.V, DisableImageGenerationChat)
+		}
+	}
+}
+
+func TestDisableImageGenerationMode_UnmarshalJSON(t *testing.T) {
+	{
+		var v DisableImageGenerationMode
+		if err := json.Unmarshal([]byte("false"), &v); err != nil {
+			t.Fatalf("unmarshal false: %v", err)
+		}
+		if v != DisableImageGenerationOff {
+			t.Fatalf("false => %v, want %v", v, DisableImageGenerationOff)
+		}
+	}
+
+	{
+		var v DisableImageGenerationMode
+		if err := json.Unmarshal([]byte("true"), &v); err != nil {
+			t.Fatalf("unmarshal true: %v", err)
+		}
+		if v != DisableImageGenerationAll {
+			t.Fatalf("true => %v, want %v", v, DisableImageGenerationAll)
+		}
+	}
+
+	{
+		var v DisableImageGenerationMode
+		if err := json.Unmarshal([]byte(`"chat"`), &v); err != nil {
+			t.Fatalf("unmarshal chat: %v", err)
+		}
+		if v != DisableImageGenerationChat {
+			t.Fatalf("chat => %v, want %v", v, DisableImageGenerationChat)
+		}
+	}
+}
@@ -9,11 +9,15 @@ type SDKConfig struct {
 	// ProxyURL is the URL of an optional proxy server to use for outbound requests.
 	ProxyURL string `yaml:"proxy-url" json:"proxy-url"`

-	// DisableImageGeneration disables the built-in image_generation tool when true.
-	// When enabled, the server will avoid injecting image_generation into request payloads,
-	// will remove any existing image_generation tool entries from tools arrays, and will
-	// return 404 for /v1/images/generations and /v1/images/edits.
-	DisableImageGeneration bool `yaml:"disable-image-generation" json:"disable-image-generation"`
+	// DisableImageGeneration controls whether the built-in image_generation tool is injected/allowed.
+	//
+	// Supported values:
+	//   - false (default): image_generation is enabled everywhere (normal behavior).
+	//   - true: image_generation is disabled everywhere. The server stops injecting it, removes it from request payloads,
+	//     and returns 404 for /v1/images/generations and /v1/images/edits.
+	//   - "chat": disable image_generation injection for all non-images endpoints (e.g. /v1/responses, /v1/chat/completions),
+	//     while keeping /v1/images/generations and /v1/images/edits enabled and preserving image_generation there.
+	DisableImageGeneration DisableImageGenerationMode `yaml:"disable-image-generation" json:"disable-image-generation"`

 	// EnableGeminiCLIEndpoint controls whether Gemini CLI internal endpoints (/v1internal:*) are enabled.
 	// Default is false for safety; when false, /v1internal:* requests are rejected.