[code-simplifier] refactor: extract resolveProxyContainerImage helper in compiler_difc_proxy (#25419) (#25435)

Author: github-actions[bot]

pkg/workflow/compiler_difc_proxy.go

@@ -201,6 +201,16 @@ func getDIFCProxyPolicyJSON(githubTool any) string {
 	return string(jsonBytes)
 }
 
+// resolveProxyContainerImage returns the full container image reference (container:version)
+// for the DIFC/CLI proxy, falling back to the default MCP gateway version if none is configured.
+func resolveProxyContainerImage(gatewayConfig *MCPGatewayRuntimeConfig) string {
+	version := gatewayConfig.Version
+	if version == "" {
+		version = string(constants.DefaultMCPGatewayVersion)
+	}
+	return gatewayConfig.Container + ":" + version
+}
+
 // buildStartDIFCProxyStepYAML returns the YAML for the "Start DIFC proxy" step,
 // or an empty string if proxy injection is not needed or the policy cannot be built.
 // This is the shared implementation used by both the main job and the indexing job.
@@ -223,14 +233,7 @@ func (c *Compiler) buildStartDIFCProxyStepYAML(data *WorkflowData) string {
 	// Resolve the container image from the MCP gateway configuration
 	// (proxy uses the same image as the gateway, just in "proxy" mode)
 	ensureDefaultMCPGatewayConfig(data)
-	gatewayConfig := data.SandboxConfig.MCP
-
-	containerImage := gatewayConfig.Container
-	if gatewayConfig.Version != "" {
-		containerImage += ":" + gatewayConfig.Version
-	} else {
-		containerImage += ":" + string(constants.DefaultMCPGatewayVersion)
-	}
+	containerImage := resolveProxyContainerImage(data.SandboxConfig.MCP)
 
 	var sb strings.Builder
 	sb.WriteString("      - name: Start DIFC proxy for pre-agent gh calls\n")
@@ -388,14 +391,7 @@ func (c *Compiler) buildStartCliProxyStepYAML(data *WorkflowData) string {
 
 	// Resolve the container image from the MCP gateway configuration
 	ensureDefaultMCPGatewayConfig(data)
-	gatewayConfig := data.SandboxConfig.MCP
-
-	containerImage := gatewayConfig.Container
-	if gatewayConfig.Version != "" {
-		containerImage += ":" + gatewayConfig.Version
-	} else {
-		containerImage += ":" + string(constants.DefaultMCPGatewayVersion)
-	}
+	containerImage := resolveProxyContainerImage(data.SandboxConfig.MCP)
 
 	var sb strings.Builder
 	sb.WriteString("      - name: Start CLI proxy\n")

pkg/workflow/compiler_difc_proxy_test.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/github/gh-aw/pkg/constants"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -758,3 +759,54 @@ func TestBuildStartCliProxyStepYAML(t *testing.T) {
 		assert.Contains(t, result, "start_cli_proxy.sh", "should reference the start script")
 	})
 }
+
+// TestResolveProxyContainerImage verifies that the helper builds the correct container
+// image reference from the gateway config, falling back to the default version when
+// no version is specified.
+func TestResolveProxyContainerImage(t *testing.T) {
+	tests := []struct {
+		name     string
+		config   *MCPGatewayRuntimeConfig
+		expected string
+	}{
+		{
+			name: "uses default version when version is empty",
+			config: &MCPGatewayRuntimeConfig{
+				Container: constants.DefaultMCPGatewayContainer,
+				Version:   "",
+			},
+			expected: constants.DefaultMCPGatewayContainer + ":" + string(constants.DefaultMCPGatewayVersion),
+		},
+		{
+			name: "uses explicit version when set",
+			config: &MCPGatewayRuntimeConfig{
+				Container: constants.DefaultMCPGatewayContainer,
+				Version:   "v1.2.3",
+			},
+			expected: constants.DefaultMCPGatewayContainer + ":v1.2.3",
+		},
+		{
+			name: "custom container with default version fallback",
+			config: &MCPGatewayRuntimeConfig{
+				Container: "ghcr.io/myorg/my-proxy",
+				Version:   "",
+			},
+			expected: "ghcr.io/myorg/my-proxy:" + string(constants.DefaultMCPGatewayVersion),
+		},
+		{
+			name: "custom container with explicit version",
+			config: &MCPGatewayRuntimeConfig{
+				Container: "ghcr.io/myorg/my-proxy",
+				Version:   "latest",
+			},
+			expected: "ghcr.io/myorg/my-proxy:latest",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := resolveProxyContainerImage(tt.config)
+			assert.Equal(t, tt.expected, got, "resolveProxyContainerImage(%+v)", tt.config)
+		})
+	}
+}