From 6ba7c810a78c9afa88550a80b90c48b24e8b4852 Mon Sep 17 00:00:00 2001
From: Luis Pater <webmaster@idotorg.org>
Date: Thu, 30 Apr 2026 12:42:08 +0800
Subject: [PATCH] feat: apply image_generation filtering before payload rules

- Updated `ApplyPayloadConfigWithRoot` to prioritize `disable-image-generation` filtering before applying payload rules.
- Ensured payload overrides can explicitly re-enable `image_generation` when required.
- Added unit tests to validate `image_generation` restoration through overrides.
---
 .../runtime/executor/helps/payload_helpers.go | 17 +++++----
 ...d_helpers_disable_image_generation_test.go | 37 +++++++++++++++++++
 2 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/internal/runtime/executor/helps/payload_helpers.go b/internal/runtime/executor/helps/payload_helpers.go
index f8905ae7..d6baba27 100644
--- a/internal/runtime/executor/helps/payload_helpers.go
+++ b/internal/runtime/executor/helps/payload_helpers.go
@@ -23,6 +23,15 @@ func ApplyPayloadConfigWithRoot(cfg *config.Config, model, protocol, root string
 	}
 	out := payload
 
+	// Apply disable-image-generation filtering before payload rules so config payload
+	// overrides can explicitly re-enable image_generation when desired.
+	if cfg.DisableImageGeneration != config.DisableImageGenerationOff {
+		if cfg.DisableImageGeneration != config.DisableImageGenerationChat || !isImagesEndpointRequestPath(requestPath) {
+			out = removeToolTypeFromPayloadWithRoot(out, root, "image_generation")
+			out = removeToolChoiceFromPayloadWithRoot(out, root, "image_generation")
+		}
+	}
+
 	rules := cfg.Payload
 	hasPayloadRules := len(rules.Default) != 0 || len(rules.DefaultRaw) != 0 || len(rules.Override) != 0 || len(rules.OverrideRaw) != 0 || len(rules.Filter) != 0
 	if hasPayloadRules {
@@ -149,14 +158,6 @@ func ApplyPayloadConfigWithRoot(cfg *config.Config, model, protocol, root string
 			}
 		}
 	}
-
-	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
 }
 
diff --git a/internal/runtime/executor/helps/payload_helpers_disable_image_generation_test.go b/internal/runtime/executor/helps/payload_helpers_disable_image_generation_test.go
index 1458d229..6fd3a0e0 100644
--- a/internal/runtime/executor/helps/payload_helpers_disable_image_generation_test.go
+++ b/internal/runtime/executor/helps/payload_helpers_disable_image_generation_test.go
@@ -95,3 +95,40 @@ func TestApplyPayloadConfigWithRoot_DisableImageGenerationChat_KeepsImageGenerat
 		t.Fatalf("expected tool_choice to be kept on images endpoint")
 	}
 }
+
+func TestApplyPayloadConfigWithRoot_DisableImageGeneration_PayloadOverrideCanRestoreImageGeneration(t *testing.T) {
+	cfg := &config.Config{
+		SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll},
+		Payload: config.PayloadConfig{
+			OverrideRaw: []config.PayloadRule{
+				{
+					Models: []config.PayloadModelRule{
+						{Name: "gpt-5.4", Protocol: "openai-response"},
+					},
+					Params: map[string]any{
+						"tools":       `[{"type":"image_generation"},{"type":"function","name":"f1"}]`,
+						"tool_choice": `{"type":"image_generation"}`,
+					},
+				},
+			},
+		},
+	}
+	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, "", "")
+
+	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 after payload override, got %d", len(arr))
+	}
+	if got := arr[0].Get("type").String(); got != "image_generation" {
+		t.Fatalf("expected first tool type=image_generation, got %q", got)
+	}
+	if !gjson.GetBytes(out, "tool_choice").Exists() {
+		t.Fatalf("expected tool_choice to be restored by payload override")
+	}
+}