mirror of
https://mirror.skon.top/github.com/router-for-me/CLIProxyAPI
synced 2026-05-01 00:30:55 +08:00
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:
@@ -1014,7 +1014,7 @@ func (s *Server) UpdateClients(cfg *config.Config) {
|
||||
}
|
||||
|
||||
if oldCfg != nil && oldCfg.DisableImageGeneration != cfg.DisableImageGeneration {
|
||||
log.Infof("disable-image-generation updated: %t -> %t", oldCfg.DisableImageGeneration, cfg.DisableImageGeneration)
|
||||
log.Infof("disable-image-generation updated: %v -> %v", oldCfg.DisableImageGeneration, cfg.DisableImageGeneration)
|
||||
}
|
||||
|
||||
applySignatureCacheConfig(oldCfg, cfg)
|
||||
|
||||
@@ -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
|
||||
|
||||
136
internal/config/disable_image_generation_mode.go
Normal file
136
internal/config/disable_image_generation_mode.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
76
internal/config/disable_image_generation_mode_test.go
Normal file
76
internal/config/disable_image_generation_mode_test.go
Normal file
@@ -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.
|
||||
|
||||
@@ -428,7 +428,8 @@ func (e *AIStudioExecutor) translateRequest(req cliproxyexecutor.Request, opts c
|
||||
}
|
||||
payload = fixGeminiImageAspectRatio(baseModel, payload)
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
payload = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", payload, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
payload = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", payload, originalTranslated, requestedModel, requestPath)
|
||||
payload, _ = sjson.DeleteBytes(payload, "generationConfig.maxOutputTokens")
|
||||
payload, _ = sjson.DeleteBytes(payload, "generationConfig.responseMimeType")
|
||||
payload, _ = sjson.DeleteBytes(payload, "generationConfig.responseJsonSchema")
|
||||
|
||||
@@ -521,7 +521,8 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
|
||||
}
|
||||
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel, requestPath)
|
||||
|
||||
useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg)
|
||||
|
||||
@@ -718,7 +719,8 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *
|
||||
}
|
||||
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel, requestPath)
|
||||
|
||||
useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg)
|
||||
|
||||
@@ -1178,7 +1180,8 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
|
||||
}
|
||||
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel, requestPath)
|
||||
|
||||
useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg)
|
||||
|
||||
|
||||
@@ -164,7 +164,8 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
|
||||
body = applyCloaking(ctx, e.cfg, auth, body, baseModel, apiKey)
|
||||
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body = ensureModelMaxTokens(body, baseModel)
|
||||
|
||||
// Disable thinking if tool_choice forces tool use (Anthropic API constraint)
|
||||
@@ -349,7 +350,8 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
|
||||
body = applyCloaking(ctx, e.cfg, auth, body, baseModel, apiKey)
|
||||
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body = ensureModelMaxTokens(body, baseModel)
|
||||
|
||||
// Disable thinking if tool_choice forces tool use (Anthropic API constraint)
|
||||
|
||||
@@ -173,7 +173,8 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re
|
||||
}
|
||||
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body, _ = sjson.SetBytes(body, "model", baseModel)
|
||||
body, _ = sjson.SetBytes(body, "stream", true)
|
||||
body, _ = sjson.DeleteBytes(body, "previous_response_id")
|
||||
@@ -181,7 +182,7 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re
|
||||
body, _ = sjson.DeleteBytes(body, "safety_identifier")
|
||||
body, _ = sjson.DeleteBytes(body, "stream_options")
|
||||
body = normalizeCodexInstructions(body)
|
||||
if e.cfg == nil || !e.cfg.DisableImageGeneration {
|
||||
if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff {
|
||||
body = ensureImageGenerationTool(body, baseModel, auth)
|
||||
}
|
||||
|
||||
@@ -327,11 +328,12 @@ func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.A
|
||||
}
|
||||
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body, _ = sjson.SetBytes(body, "model", baseModel)
|
||||
body, _ = sjson.DeleteBytes(body, "stream")
|
||||
body = normalizeCodexInstructions(body)
|
||||
if e.cfg == nil || !e.cfg.DisableImageGeneration {
|
||||
if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff {
|
||||
body = ensureImageGenerationTool(body, baseModel, auth)
|
||||
}
|
||||
|
||||
@@ -421,14 +423,15 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
|
||||
}
|
||||
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body, _ = sjson.DeleteBytes(body, "previous_response_id")
|
||||
body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")
|
||||
body, _ = sjson.DeleteBytes(body, "safety_identifier")
|
||||
body, _ = sjson.DeleteBytes(body, "stream_options")
|
||||
body, _ = sjson.SetBytes(body, "model", baseModel)
|
||||
body = normalizeCodexInstructions(body)
|
||||
if e.cfg == nil || !e.cfg.DisableImageGeneration {
|
||||
if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff {
|
||||
body = ensureImageGenerationTool(body, baseModel, auth)
|
||||
}
|
||||
|
||||
|
||||
@@ -184,14 +184,16 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut
|
||||
}
|
||||
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body, _ = sjson.SetBytes(body, "model", baseModel)
|
||||
body, _ = sjson.SetBytes(body, "stream", true)
|
||||
body, _ = sjson.DeleteBytes(body, "previous_response_id")
|
||||
body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")
|
||||
body, _ = sjson.DeleteBytes(body, "safety_identifier")
|
||||
if !gjson.GetBytes(body, "instructions").Exists() {
|
||||
body, _ = sjson.SetBytes(body, "instructions", "")
|
||||
body = normalizeCodexInstructions(body)
|
||||
if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff {
|
||||
body = ensureImageGenerationTool(body, baseModel, auth)
|
||||
}
|
||||
|
||||
httpURL := strings.TrimSuffix(baseURL, "/") + "/responses"
|
||||
@@ -387,7 +389,12 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
|
||||
}
|
||||
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, body, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, body, requestedModel, requestPath)
|
||||
body = normalizeCodexInstructions(body)
|
||||
if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff {
|
||||
body = ensureImageGenerationTool(body, baseModel, auth)
|
||||
}
|
||||
|
||||
httpURL := strings.TrimSuffix(baseURL, "/") + "/responses"
|
||||
wsURL, err := buildCodexResponsesWebsocketURL(httpURL)
|
||||
|
||||
@@ -139,7 +139,8 @@ func (e *GeminiCLIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth
|
||||
|
||||
basePayload = fixGeminiCLIImageAspectRatio(baseModel, basePayload)
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
basePayload = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "gemini", "request", basePayload, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
basePayload = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "gemini", "request", basePayload, originalTranslated, requestedModel, requestPath)
|
||||
|
||||
action := "generateContent"
|
||||
if req.Metadata != nil {
|
||||
@@ -294,7 +295,8 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
|
||||
|
||||
basePayload = fixGeminiCLIImageAspectRatio(baseModel, basePayload)
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
basePayload = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "gemini", "request", basePayload, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
basePayload = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "gemini", "request", basePayload, originalTranslated, requestedModel, requestPath)
|
||||
|
||||
projectID := resolveGeminiProjectID(auth)
|
||||
|
||||
|
||||
@@ -132,7 +132,8 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
|
||||
|
||||
body = fixGeminiImageAspectRatio(baseModel, body)
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body, _ = sjson.SetBytes(body, "model", baseModel)
|
||||
|
||||
action := "generateContent"
|
||||
@@ -239,7 +240,8 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
|
||||
|
||||
body = fixGeminiImageAspectRatio(baseModel, body)
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body, _ = sjson.SetBytes(body, "model", baseModel)
|
||||
|
||||
baseURL := resolveGeminiBaseURL(auth)
|
||||
|
||||
@@ -335,7 +335,8 @@ func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, au
|
||||
|
||||
body = fixGeminiImageAspectRatio(baseModel, body)
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body, _ = sjson.SetBytes(body, "model", baseModel)
|
||||
}
|
||||
|
||||
@@ -455,7 +456,8 @@ func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *clip
|
||||
|
||||
body = fixGeminiImageAspectRatio(baseModel, body)
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body, _ = sjson.SetBytes(body, "model", baseModel)
|
||||
|
||||
action := getVertexAction(baseModel, false)
|
||||
@@ -565,7 +567,8 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte
|
||||
|
||||
body = fixGeminiImageAspectRatio(baseModel, body)
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body, _ = sjson.SetBytes(body, "model", baseModel)
|
||||
|
||||
action := getVertexAction(baseModel, true)
|
||||
@@ -694,7 +697,8 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth
|
||||
|
||||
body = fixGeminiImageAspectRatio(baseModel, body)
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body, _ = sjson.SetBytes(body, "model", baseModel)
|
||||
|
||||
action := getVertexAction(baseModel, true)
|
||||
|
||||
@@ -16,7 +16,8 @@ import (
|
||||
// and restricts matches to the given protocol when supplied. Defaults are checked
|
||||
// against the original payload when provided. requestedModel carries the client-visible
|
||||
// model name before alias resolution so payload rules can target aliases precisely.
|
||||
func ApplyPayloadConfigWithRoot(cfg *config.Config, model, protocol, root string, payload, original []byte, requestedModel string) []byte {
|
||||
// requestPath is the inbound HTTP request path (when available) used for endpoint-scoped gates.
|
||||
func ApplyPayloadConfigWithRoot(cfg *config.Config, model, protocol, root string, payload, original []byte, requestedModel string, requestPath string) []byte {
|
||||
if cfg == nil || len(payload) == 0 {
|
||||
return payload
|
||||
}
|
||||
@@ -149,13 +150,34 @@ func ApplyPayloadConfigWithRoot(cfg *config.Config, model, protocol, root string
|
||||
}
|
||||
}
|
||||
|
||||
if cfg.DisableImageGeneration {
|
||||
if cfg.DisableImageGeneration != config.DisableImageGenerationOff {
|
||||
if cfg.DisableImageGeneration == config.DisableImageGenerationChat && isImagesEndpointRequestPath(requestPath) {
|
||||
return out
|
||||
}
|
||||
out = removeToolTypeFromPayloadWithRoot(out, root, "image_generation")
|
||||
out = removeToolChoiceFromPayloadWithRoot(out, root, "image_generation")
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func isImagesEndpointRequestPath(path string) bool {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return false
|
||||
}
|
||||
if path == "/v1/images/generations" || path == "/v1/images/edits" {
|
||||
return true
|
||||
}
|
||||
// Be tolerant of prefix routers that may report a longer matched route.
|
||||
if strings.HasSuffix(path, "/v1/images/generations") || strings.HasSuffix(path, "/v1/images/edits") {
|
||||
return true
|
||||
}
|
||||
if strings.HasSuffix(path, "/images/generations") || strings.HasSuffix(path, "/images/edits") {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func payloadModelRulesMatch(rules []config.PayloadModelRule, protocol string, models []string) bool {
|
||||
if len(rules) == 0 || len(models) == 0 {
|
||||
return false
|
||||
@@ -367,6 +389,24 @@ func PayloadRequestedModel(opts cliproxyexecutor.Options, fallback string) strin
|
||||
}
|
||||
}
|
||||
|
||||
func PayloadRequestPath(opts cliproxyexecutor.Options) string {
|
||||
if len(opts.Metadata) == 0 {
|
||||
return ""
|
||||
}
|
||||
raw, ok := opts.Metadata[cliproxyexecutor.RequestPathMetadataKey]
|
||||
if !ok || raw == nil {
|
||||
return ""
|
||||
}
|
||||
switch v := raw.(type) {
|
||||
case string:
|
||||
return strings.TrimSpace(v)
|
||||
case []byte:
|
||||
return strings.TrimSpace(string(v))
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// matchModelPattern performs simple wildcard matching where '*' matches zero or more characters.
|
||||
// Examples:
|
||||
//
|
||||
|
||||
@@ -9,11 +9,11 @@ import (
|
||||
|
||||
func TestApplyPayloadConfigWithRoot_DisableImageGeneration_RemovesToolsEntry(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
SDKConfig: config.SDKConfig{DisableImageGeneration: true},
|
||||
SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll},
|
||||
}
|
||||
payload := []byte(`{"tools":[{"type":"image_generation","output_format":"png"},{"type":"function","name":"f1"}]}`)
|
||||
|
||||
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "openai-response", "", payload, nil, "")
|
||||
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "openai-response", "", payload, nil, "", "")
|
||||
|
||||
tools := gjson.GetBytes(out, "tools")
|
||||
if !tools.Exists() || !tools.IsArray() {
|
||||
@@ -30,11 +30,11 @@ func TestApplyPayloadConfigWithRoot_DisableImageGeneration_RemovesToolsEntry(t *
|
||||
|
||||
func TestApplyPayloadConfigWithRoot_DisableImageGeneration_RemovesToolsEntryWithRoot(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
SDKConfig: config.SDKConfig{DisableImageGeneration: true},
|
||||
SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll},
|
||||
}
|
||||
payload := []byte(`{"request":{"tools":[{"type":"image_generation"},{"type":"web_search"}]}}`)
|
||||
|
||||
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "gemini-cli", "request", payload, nil, "")
|
||||
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "gemini-cli", "request", payload, nil, "", "")
|
||||
|
||||
tools := gjson.GetBytes(out, "request.tools")
|
||||
if !tools.Exists() || !tools.IsArray() {
|
||||
@@ -51,11 +51,11 @@ func TestApplyPayloadConfigWithRoot_DisableImageGeneration_RemovesToolsEntryWith
|
||||
|
||||
func TestApplyPayloadConfigWithRoot_DisableImageGeneration_RemovesToolChoiceByType(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
SDKConfig: config.SDKConfig{DisableImageGeneration: true},
|
||||
SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll},
|
||||
}
|
||||
payload := []byte(`{"tools":[{"type":"image_generation"},{"type":"function","name":"f1"}],"tool_choice":{"type":"image_generation"}}`)
|
||||
|
||||
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "openai-response", "", payload, nil, "")
|
||||
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "openai-response", "", payload, nil, "", "")
|
||||
|
||||
if gjson.GetBytes(out, "tool_choice").Exists() {
|
||||
t.Fatalf("expected tool_choice to be removed")
|
||||
@@ -64,13 +64,34 @@ func TestApplyPayloadConfigWithRoot_DisableImageGeneration_RemovesToolChoiceByTy
|
||||
|
||||
func TestApplyPayloadConfigWithRoot_DisableImageGeneration_RemovesToolChoiceByNameWithRoot(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
SDKConfig: config.SDKConfig{DisableImageGeneration: true},
|
||||
SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll},
|
||||
}
|
||||
payload := []byte(`{"request":{"tools":[{"type":"image_generation"},{"type":"web_search"}],"tool_choice":{"type":"tool","name":"image_generation"}}}`)
|
||||
|
||||
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "gemini-cli", "request", payload, nil, "")
|
||||
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "gemini-cli", "request", payload, nil, "", "")
|
||||
|
||||
if gjson.GetBytes(out, "request.tool_choice").Exists() {
|
||||
t.Fatalf("expected request.tool_choice to be removed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyPayloadConfigWithRoot_DisableImageGenerationChat_KeepsImageGenerationOnImagesEndpoints(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationChat},
|
||||
}
|
||||
payload := []byte(`{"tools":[{"type":"image_generation"},{"type":"function","name":"f1"}],"tool_choice":{"type":"image_generation"}}`)
|
||||
|
||||
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "openai-response", "", payload, nil, "", "/v1/images/generations")
|
||||
|
||||
tools := gjson.GetBytes(out, "tools")
|
||||
if !tools.Exists() || !tools.IsArray() {
|
||||
t.Fatalf("expected tools array, got %v", tools.Type)
|
||||
}
|
||||
arr := tools.Array()
|
||||
if len(arr) != 2 {
|
||||
t.Fatalf("expected 2 tools (no removal), got %d", len(arr))
|
||||
}
|
||||
if !gjson.GetBytes(out, "tool_choice").Exists() {
|
||||
t.Fatalf("expected tool_choice to be kept on images endpoint")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -108,7 +108,8 @@ func (e *KimiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req
|
||||
}
|
||||
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body, err = normalizeKimiToolMessageLinks(body)
|
||||
if err != nil {
|
||||
return resp, err
|
||||
@@ -217,7 +218,8 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut
|
||||
return nil, fmt.Errorf("kimi executor: failed to set stream_options in payload: %w", err)
|
||||
}
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
body = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel, requestPath)
|
||||
body, err = normalizeKimiToolMessageLinks(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -97,7 +97,8 @@ func (e *OpenAICompatExecutor) Execute(ctx context.Context, auth *cliproxyauth.A
|
||||
originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, opts.Stream)
|
||||
translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, opts.Stream)
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", translated, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", translated, originalTranslated, requestedModel, requestPath)
|
||||
if opts.Alt == "responses/compact" {
|
||||
if updated, errDelete := sjson.DeleteBytes(translated, "stream"); errDelete == nil {
|
||||
translated = updated
|
||||
@@ -199,7 +200,8 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy
|
||||
originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)
|
||||
translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true)
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", translated, originalTranslated, requestedModel)
|
||||
requestPath := helps.PayloadRequestPath(opts)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", translated, originalTranslated, requestedModel, requestPath)
|
||||
|
||||
translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier())
|
||||
if err != nil {
|
||||
|
||||
@@ -43,7 +43,7 @@ func BuildConfigChangeDetails(oldCfg, newCfg *config.Config) []string {
|
||||
changes = append(changes, fmt.Sprintf("disable-cooling: %t -> %t", oldCfg.DisableCooling, newCfg.DisableCooling))
|
||||
}
|
||||
if oldCfg.DisableImageGeneration != newCfg.DisableImageGeneration {
|
||||
changes = append(changes, fmt.Sprintf("disable-image-generation: %t -> %t", oldCfg.DisableImageGeneration, newCfg.DisableImageGeneration))
|
||||
changes = append(changes, fmt.Sprintf("disable-image-generation: %v -> %v", oldCfg.DisableImageGeneration, newCfg.DisableImageGeneration))
|
||||
}
|
||||
if oldCfg.RequestLog != newCfg.RequestLog {
|
||||
changes = append(changes, fmt.Sprintf("request-log: %t -> %t", oldCfg.RequestLog, newCfg.RequestLog))
|
||||
|
||||
@@ -279,7 +279,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) {
|
||||
APIKeys: []string{" key-1 ", "key-2"},
|
||||
ForceModelPrefix: true,
|
||||
NonStreamKeepAliveInterval: 5,
|
||||
DisableImageGeneration: true,
|
||||
DisableImageGeneration: config.DisableImageGenerationAll,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -408,7 +408,7 @@ func TestBuildConfigChangeDetails_AllBranches(t *testing.T) {
|
||||
RequestLog: true,
|
||||
ProxyURL: "http://new-proxy",
|
||||
APIKeys: []string{"keyB"},
|
||||
DisableImageGeneration: true,
|
||||
DisableImageGeneration: config.DisableImageGenerationAll,
|
||||
},
|
||||
OAuthExcludedModels: map[string][]string{"p1": {"b", "c"}, "p2": {"d"}},
|
||||
OpenAICompatibility: []config.OpenAICompatibility{
|
||||
|
||||
Reference in New Issue
Block a user