Skip to content

Commit e3e60f9

Browse files
committed
feat: support disabling image generation globally
- Added `disable-image-generation` configuration flag to disable the `image_generation` tool globally. - Updated payload handling to remove `image_generation` tools from request payload arrays when the flag is enabled. - Modified OpenAI image handlers (`ImagesGenerations`, `ImagesEdits`) to return 404 when the feature is disabled. - Enhanced configuration diff logging to track changes for the `disable-image-generation` flag. - Added accompanying unit tests for the new feature in payload helpers and image handler logic.
1 parent 359ec30 commit e3e60f9

11 files changed

Lines changed: 282 additions & 124 deletions

File tree

config.example.yaml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,10 @@ max-retry-interval: 30
9090
# When true, disable auth/model cooldown scheduling globally (prevents blackout windows after failure states).
9191
disable-cooling: false
9292

93+
# When true, disable the built-in image_generation tool globally.
94+
# The server will stop injecting image_generation and will also remove it from request payload tools arrays.
95+
disable-image-generation: false
96+
9397
# Core auth auto-refresh worker pool size (OAuth/file-based auth token refresh).
9498
# When > 0, overrides the default worker count (16).
9599
# auth-auto-refresh-workers: 16

internal/api/server.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1013,6 +1013,10 @@ func (s *Server) UpdateClients(cfg *config.Config) {
10131013
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
10141014
}
10151015

1016+
if oldCfg != nil && oldCfg.DisableImageGeneration != cfg.DisableImageGeneration {
1017+
log.Infof("disable-image-generation updated: %t -> %t", oldCfg.DisableImageGeneration, cfg.DisableImageGeneration)
1018+
}
1019+
10161020
applySignatureCacheConfig(oldCfg, cfg)
10171021

10181022
if s.handlers != nil && s.handlers.AuthManager != nil {

internal/config/config.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -610,6 +610,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
610610
cfg.ErrorLogsMaxFiles = 10
611611
cfg.UsageStatisticsEnabled = false
612612
cfg.DisableCooling = false
613+
cfg.DisableImageGeneration = false
613614
cfg.Pprof.Enable = false
614615
cfg.Pprof.Addr = DefaultPprofAddr
615616
cfg.AmpCode.RestrictManagementToLocalhost = false // Default to false: API key auth is sufficient

internal/config/sdk_config.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,12 @@ type SDKConfig struct {
99
// ProxyURL is the URL of an optional proxy server to use for outbound requests.
1010
ProxyURL string `yaml:"proxy-url" json:"proxy-url"`
1111

12+
// DisableImageGeneration disables the built-in image_generation tool when true.
13+
// When enabled, the server will avoid injecting image_generation into request payloads,
14+
// will remove any existing image_generation tool entries from tools arrays, and will
15+
// return 404 for /v1/images/generations and /v1/images/edits.
16+
DisableImageGeneration bool `yaml:"disable-image-generation" json:"disable-image-generation"`
17+
1218
// EnableGeminiCLIEndpoint controls whether Gemini CLI internal endpoints (/v1internal:*) are enabled.
1319
// Default is false for safety; when false, /v1internal:* requests are rejected.
1420
EnableGeminiCLIEndpoint bool `yaml:"enable-gemini-cli-endpoint" json:"enable-gemini-cli-endpoint"`

internal/runtime/executor/codex_executor.go

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,9 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re
181181
body, _ = sjson.DeleteBytes(body, "safety_identifier")
182182
body, _ = sjson.DeleteBytes(body, "stream_options")
183183
body = normalizeCodexInstructions(body)
184-
body = ensureImageGenerationTool(body, baseModel, auth)
184+
if e.cfg == nil || !e.cfg.DisableImageGeneration {
185+
body = ensureImageGenerationTool(body, baseModel, auth)
186+
}
185187

186188
url := strings.TrimSuffix(baseURL, "/") + "/responses"
187189
httpReq, err := e.cacheHelper(ctx, from, url, req, body)
@@ -329,7 +331,9 @@ func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.A
329331
body, _ = sjson.SetBytes(body, "model", baseModel)
330332
body, _ = sjson.DeleteBytes(body, "stream")
331333
body = normalizeCodexInstructions(body)
332-
body = ensureImageGenerationTool(body, baseModel, auth)
334+
if e.cfg == nil || !e.cfg.DisableImageGeneration {
335+
body = ensureImageGenerationTool(body, baseModel, auth)
336+
}
333337

334338
url := strings.TrimSuffix(baseURL, "/") + "/responses/compact"
335339
httpReq, err := e.cacheHelper(ctx, from, url, req, body)
@@ -424,7 +428,9 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
424428
body, _ = sjson.DeleteBytes(body, "stream_options")
425429
body, _ = sjson.SetBytes(body, "model", baseModel)
426430
body = normalizeCodexInstructions(body)
427-
body = ensureImageGenerationTool(body, baseModel, auth)
431+
if e.cfg == nil || !e.cfg.DisableImageGeneration {
432+
body = ensureImageGenerationTool(body, baseModel, auth)
433+
}
428434

429435
url := strings.TrimSuffix(baseURL, "/") + "/responses"
430436
httpReq, err := e.cacheHelper(ctx, from, url, req, body)

internal/runtime/executor/helps/payload_helpers.go

Lines changed: 162 additions & 118 deletions
Original file line numberDiff line numberDiff line change
@@ -20,133 +20,137 @@ func ApplyPayloadConfigWithRoot(cfg *config.Config, model, protocol, root string
2020
if cfg == nil || len(payload) == 0 {
2121
return payload
2222
}
23-
rules := cfg.Payload
24-
if len(rules.Default) == 0 && len(rules.DefaultRaw) == 0 && len(rules.Override) == 0 && len(rules.OverrideRaw) == 0 && len(rules.Filter) == 0 {
25-
return payload
26-
}
27-
model = strings.TrimSpace(model)
28-
requestedModel = strings.TrimSpace(requestedModel)
29-
if model == "" && requestedModel == "" {
30-
return payload
31-
}
32-
candidates := payloadModelCandidates(model, requestedModel)
3323
out := payload
34-
source := original
35-
if len(source) == 0 {
36-
source = payload
37-
}
38-
appliedDefaults := make(map[string]struct{})
39-
// Apply default rules: first write wins per field across all matching rules.
40-
for i := range rules.Default {
41-
rule := &rules.Default[i]
42-
if !payloadModelRulesMatch(rule.Models, protocol, candidates) {
43-
continue
44-
}
45-
for path, value := range rule.Params {
46-
fullPath := buildPayloadPath(root, path)
47-
if fullPath == "" {
48-
continue
49-
}
50-
if gjson.GetBytes(source, fullPath).Exists() {
51-
continue
52-
}
53-
if _, ok := appliedDefaults[fullPath]; ok {
54-
continue
55-
}
56-
updated, errSet := sjson.SetBytes(out, fullPath, value)
57-
if errSet != nil {
58-
continue
59-
}
60-
out = updated
61-
appliedDefaults[fullPath] = struct{}{}
62-
}
63-
}
64-
// Apply default raw rules: first write wins per field across all matching rules.
65-
for i := range rules.DefaultRaw {
66-
rule := &rules.DefaultRaw[i]
67-
if !payloadModelRulesMatch(rule.Models, protocol, candidates) {
68-
continue
69-
}
70-
for path, value := range rule.Params {
71-
fullPath := buildPayloadPath(root, path)
72-
if fullPath == "" {
73-
continue
74-
}
75-
if gjson.GetBytes(source, fullPath).Exists() {
76-
continue
77-
}
78-
if _, ok := appliedDefaults[fullPath]; ok {
79-
continue
80-
}
81-
rawValue, ok := payloadRawValue(value)
82-
if !ok {
83-
continue
84-
}
85-
updated, errSet := sjson.SetRawBytes(out, fullPath, rawValue)
86-
if errSet != nil {
87-
continue
24+
25+
rules := cfg.Payload
26+
hasPayloadRules := len(rules.Default) != 0 || len(rules.DefaultRaw) != 0 || len(rules.Override) != 0 || len(rules.OverrideRaw) != 0 || len(rules.Filter) != 0
27+
if hasPayloadRules {
28+
model = strings.TrimSpace(model)
29+
requestedModel = strings.TrimSpace(requestedModel)
30+
if model != "" || requestedModel != "" {
31+
candidates := payloadModelCandidates(model, requestedModel)
32+
source := original
33+
if len(source) == 0 {
34+
source = payload
8835
}
89-
out = updated
90-
appliedDefaults[fullPath] = struct{}{}
91-
}
92-
}
93-
// Apply override rules: last write wins per field across all matching rules.
94-
for i := range rules.Override {
95-
rule := &rules.Override[i]
96-
if !payloadModelRulesMatch(rule.Models, protocol, candidates) {
97-
continue
98-
}
99-
for path, value := range rule.Params {
100-
fullPath := buildPayloadPath(root, path)
101-
if fullPath == "" {
102-
continue
36+
appliedDefaults := make(map[string]struct{})
37+
// Apply default rules: first write wins per field across all matching rules.
38+
for i := range rules.Default {
39+
rule := &rules.Default[i]
40+
if !payloadModelRulesMatch(rule.Models, protocol, candidates) {
41+
continue
42+
}
43+
for path, value := range rule.Params {
44+
fullPath := buildPayloadPath(root, path)
45+
if fullPath == "" {
46+
continue
47+
}
48+
if gjson.GetBytes(source, fullPath).Exists() {
49+
continue
50+
}
51+
if _, ok := appliedDefaults[fullPath]; ok {
52+
continue
53+
}
54+
updated, errSet := sjson.SetBytes(out, fullPath, value)
55+
if errSet != nil {
56+
continue
57+
}
58+
out = updated
59+
appliedDefaults[fullPath] = struct{}{}
60+
}
10361
}
104-
updated, errSet := sjson.SetBytes(out, fullPath, value)
105-
if errSet != nil {
106-
continue
62+
// Apply default raw rules: first write wins per field across all matching rules.
63+
for i := range rules.DefaultRaw {
64+
rule := &rules.DefaultRaw[i]
65+
if !payloadModelRulesMatch(rule.Models, protocol, candidates) {
66+
continue
67+
}
68+
for path, value := range rule.Params {
69+
fullPath := buildPayloadPath(root, path)
70+
if fullPath == "" {
71+
continue
72+
}
73+
if gjson.GetBytes(source, fullPath).Exists() {
74+
continue
75+
}
76+
if _, ok := appliedDefaults[fullPath]; ok {
77+
continue
78+
}
79+
rawValue, ok := payloadRawValue(value)
80+
if !ok {
81+
continue
82+
}
83+
updated, errSet := sjson.SetRawBytes(out, fullPath, rawValue)
84+
if errSet != nil {
85+
continue
86+
}
87+
out = updated
88+
appliedDefaults[fullPath] = struct{}{}
89+
}
10790
}
108-
out = updated
109-
}
110-
}
111-
// Apply override raw rules: last write wins per field across all matching rules.
112-
for i := range rules.OverrideRaw {
113-
rule := &rules.OverrideRaw[i]
114-
if !payloadModelRulesMatch(rule.Models, protocol, candidates) {
115-
continue
116-
}
117-
for path, value := range rule.Params {
118-
fullPath := buildPayloadPath(root, path)
119-
if fullPath == "" {
120-
continue
91+
// Apply override rules: last write wins per field across all matching rules.
92+
for i := range rules.Override {
93+
rule := &rules.Override[i]
94+
if !payloadModelRulesMatch(rule.Models, protocol, candidates) {
95+
continue
96+
}
97+
for path, value := range rule.Params {
98+
fullPath := buildPayloadPath(root, path)
99+
if fullPath == "" {
100+
continue
101+
}
102+
updated, errSet := sjson.SetBytes(out, fullPath, value)
103+
if errSet != nil {
104+
continue
105+
}
106+
out = updated
107+
}
121108
}
122-
rawValue, ok := payloadRawValue(value)
123-
if !ok {
124-
continue
109+
// Apply override raw rules: last write wins per field across all matching rules.
110+
for i := range rules.OverrideRaw {
111+
rule := &rules.OverrideRaw[i]
112+
if !payloadModelRulesMatch(rule.Models, protocol, candidates) {
113+
continue
114+
}
115+
for path, value := range rule.Params {
116+
fullPath := buildPayloadPath(root, path)
117+
if fullPath == "" {
118+
continue
119+
}
120+
rawValue, ok := payloadRawValue(value)
121+
if !ok {
122+
continue
123+
}
124+
updated, errSet := sjson.SetRawBytes(out, fullPath, rawValue)
125+
if errSet != nil {
126+
continue
127+
}
128+
out = updated
129+
}
125130
}
126-
updated, errSet := sjson.SetRawBytes(out, fullPath, rawValue)
127-
if errSet != nil {
128-
continue
131+
// Apply filter rules: remove matching paths from payload.
132+
for i := range rules.Filter {
133+
rule := &rules.Filter[i]
134+
if !payloadModelRulesMatch(rule.Models, protocol, candidates) {
135+
continue
136+
}
137+
for _, path := range rule.Params {
138+
fullPath := buildPayloadPath(root, path)
139+
if fullPath == "" {
140+
continue
141+
}
142+
updated, errDel := sjson.DeleteBytes(out, fullPath)
143+
if errDel != nil {
144+
continue
145+
}
146+
out = updated
147+
}
129148
}
130-
out = updated
131149
}
132150
}
133-
// Apply filter rules: remove matching paths from payload.
134-
for i := range rules.Filter {
135-
rule := &rules.Filter[i]
136-
if !payloadModelRulesMatch(rule.Models, protocol, candidates) {
137-
continue
138-
}
139-
for _, path := range rule.Params {
140-
fullPath := buildPayloadPath(root, path)
141-
if fullPath == "" {
142-
continue
143-
}
144-
updated, errDel := sjson.DeleteBytes(out, fullPath)
145-
if errDel != nil {
146-
continue
147-
}
148-
out = updated
149-
}
151+
152+
if cfg.DisableImageGeneration {
153+
out = removeToolTypeFromPayloadWithRoot(out, root, "image_generation")
150154
}
151155
return out
152156
}
@@ -226,6 +230,46 @@ func buildPayloadPath(root, path string) string {
226230
return r + "." + p
227231
}
228232

233+
func removeToolTypeFromPayloadWithRoot(payload []byte, root string, toolType string) []byte {
234+
if len(payload) == 0 {
235+
return payload
236+
}
237+
toolType = strings.TrimSpace(toolType)
238+
if toolType == "" {
239+
return payload
240+
}
241+
toolsPath := buildPayloadPath(root, "tools")
242+
return removeToolTypeFromToolsArray(payload, toolsPath, toolType)
243+
}
244+
245+
func removeToolTypeFromToolsArray(payload []byte, toolsPath string, toolType string) []byte {
246+
tools := gjson.GetBytes(payload, toolsPath)
247+
if !tools.Exists() || !tools.IsArray() {
248+
return payload
249+
}
250+
removed := false
251+
filtered := []byte(`[]`)
252+
for _, tool := range tools.Array() {
253+
if tool.Get("type").String() == toolType {
254+
removed = true
255+
continue
256+
}
257+
updated, errSet := sjson.SetRawBytes(filtered, "-1", []byte(tool.Raw))
258+
if errSet != nil {
259+
continue
260+
}
261+
filtered = updated
262+
}
263+
if !removed {
264+
return payload
265+
}
266+
updated, errSet := sjson.SetRawBytes(payload, toolsPath, filtered)
267+
if errSet != nil {
268+
return payload
269+
}
270+
return updated
271+
}
272+
229273
func payloadRawValue(value any) ([]byte, bool) {
230274
if value == nil {
231275
return nil, false

0 commit comments

Comments
 (0)