Refactor API proxy to use supportsLLMGateway engine flag

pkg/workflow/claude_engine.go

@@ -25,12 +25,12 @@ func NewClaudeEngine() *ClaudeEngine {
 			description:            "Uses Claude Code with full MCP tool support and allow-listing",
 			experimental:           false,
 			supportsToolsAllowlist: true,
-			supportsHTTPTransport:  true,  // Claude supports both stdio and HTTP transport
-			supportsMaxTurns:       true,  // Claude supports max-turns feature
-			supportsWebFetch:       true,  // Claude has built-in WebFetch support
-			supportsWebSearch:      true,  // Claude has built-in WebSearch support
-			supportsFirewall:       true,  // Claude supports network firewalling via AWF
-			supportsLLMGateway:     false, // Claude does not support LLM gateway
+			supportsHTTPTransport:  true, // Claude supports both stdio and HTTP transport
+			supportsMaxTurns:       true, // Claude supports max-turns feature
+			supportsWebFetch:       true, // Claude has built-in WebFetch support
+			supportsWebSearch:      true, // Claude has built-in WebSearch support
+			supportsFirewall:       true, // Claude supports network firewalling via AWF
+			supportsLLMGateway:     true, // Claude supports LLM gateway via AWF api-proxy
 		},
 	}
 }
@@ -328,11 +328,12 @@ func (e *ClaudeEngine) GetExecutionSteps(workflowData *WorkflowData, logFile str
 		awfArgs = append(awfArgs, "--skip-pull")
 		claudeLog.Print("Using --skip-pull since images are pre-downloaded")
 
-		// Enable API proxy sidecar for secure credential management
-		// The api-proxy container holds the ANTHROPIC_API_KEY and proxies
-		// requests to api.anthropic.com through the firewall
-		awfArgs = append(awfArgs, "--enable-api-proxy")
-		claudeLog.Print("Added --enable-api-proxy for Claude API proxying")
+		// Enable API proxy sidecar if this engine supports LLM gateway
+		// The api-proxy container holds the LLM API keys and proxies requests through the firewall
+		if e.SupportsLLMGateway() {
+			awfArgs = append(awfArgs, "--enable-api-proxy")
+			claudeLog.Print("Added --enable-api-proxy for LLM API proxying")
+		}
 
 		// Add SSL Bump support for HTTPS content inspection (v0.9.0+)
 		sslBumpArgs := getSSLBumpArgs(firewallConfig)

pkg/workflow/codex_engine.go

@@ -41,7 +41,7 @@ func NewCodexEngine() *CodexEngine {
 			supportsWebFetch:       false, // Codex does not have built-in web-fetch support
 			supportsWebSearch:      true,  // Codex has built-in web-search support
 			supportsFirewall:       true,  // Codex supports network firewalling via AWF
-			supportsLLMGateway:     true,  // Codex supports LLM gateway
+			supportsLLMGateway:     false, // Codex does not support LLM gateway
 		},
 	}
 }
@@ -247,6 +247,13 @@ func (e *CodexEngine) GetExecutionSteps(workflowData *WorkflowData, logFile stri
 		awfArgs = append(awfArgs, "--skip-pull")
 		codexEngineLog.Print("Using --skip-pull since images are pre-downloaded")
 
+		// Enable API proxy sidecar if this engine supports LLM gateway
+		// The api-proxy container holds the LLM API keys and proxies requests through the firewall
+		if e.SupportsLLMGateway() {
+			awfArgs = append(awfArgs, "--enable-api-proxy")
+			codexEngineLog.Print("Added --enable-api-proxy for LLM API proxying")
+		}
+
 		// Note: No --tty flag for Codex (it's not a TUI, it outputs to stdout/stderr)
 
 		// Add SSL Bump support for HTTPS content inspection (v0.9.0+)

pkg/workflow/docker.go

@@ -83,7 +83,7 @@ func collectDockerImages(tools map[string]any, workflowData *WorkflowData, actio
 	}
 
 	// Collect AWF (firewall) container images when firewall is enabled
-	// AWF uses three containers: squid (proxy), agent, and api-proxy (for Claude/Codex)
+	// AWF uses three containers: squid (proxy), agent, and api-proxy (for engines with LLM gateway support)
 	if isFirewallEnabled(workflowData) {
 		// Get the firewall version for image tags
 		firewallConfig := getFirewallConfig(workflowData)
@@ -105,15 +105,21 @@ func collectDockerImages(tools map[string]any, workflowData *WorkflowData, actio
 			dockerLog.Printf("Added AWF agent container: %s", agentImage)
 		}
 
-		// Add api-proxy sidecar container for engines that use --enable-api-proxy
+		// Add api-proxy sidecar container for engines that support LLM gateway
 		// The api-proxy holds LLM API keys securely and proxies requests through Squid:
+		//   - Port 10000: OpenAI API proxy (for Codex)
 		//   - Port 10001: Anthropic API proxy (for Claude)
-		if workflowData != nil && workflowData.AI == "claude" {
-			apiProxyImage := constants.DefaultFirewallRegistry + "/api-proxy:" + awfImageTag
-			if !imageSet[apiProxyImage] {
-				images = append(images, apiProxyImage)
-				imageSet[apiProxyImage] = true
-				dockerLog.Printf("Added AWF api-proxy sidecar container: %s", apiProxyImage)
+		// Check if the engine supports LLM gateway by querying the engine registry
+		if workflowData != nil && workflowData.AI != "" {
+			registry := GetGlobalEngineRegistry()
+			engine, err := registry.GetEngine(workflowData.AI)
+			if err == nil && engine.SupportsLLMGateway() {
+				apiProxyImage := constants.DefaultFirewallRegistry + "/api-proxy:" + awfImageTag
+				if !imageSet[apiProxyImage] {
+					images = append(images, apiProxyImage)
+					imageSet[apiProxyImage] = true
+					dockerLog.Printf("Added AWF api-proxy sidecar container for engine with LLM gateway support: %s", apiProxyImage)
+				}
 			}
 		}
 	}

pkg/workflow/docker_api_proxy_test.go

@@ -6,7 +6,7 @@ import (
 	"github.com/github/gh-aw/pkg/constants"
 )
 
-func TestCollectDockerImages_APIProxyForClaude(t *testing.T) {
+func TestCollectDockerImages_APIProxyForEnginesWithLLMGateway(t *testing.T) {
 	awfImageTag := "0.16.5"
 
 	tests := []struct {
@@ -15,15 +15,20 @@ func TestCollectDockerImages_APIProxyForClaude(t *testing.T) {
 		expectAPIProxy bool
 	}{
 		{
-			name:           "Claude engine includes api-proxy image",
+			name:           "Claude engine includes api-proxy image (supportsLLMGateway: true)",
 			engine:         "claude",
 			expectAPIProxy: true,
 		},
 		{
-			name:           "Copilot engine does not include api-proxy image",
+			name:           "Copilot engine does not include api-proxy image (supportsLLMGateway: false)",
 			engine:         "copilot",
 			expectAPIProxy: false,
 		},
+		{
+			name:           "Codex engine does not include api-proxy image (supportsLLMGateway: false)",
+			engine:         "codex",
+			expectAPIProxy: false,
+		},
 	}
 
 	for _, tt := range tests {

pkg/workflow/enable_api_proxy_test.go

@@ -5,10 +5,10 @@ import (
 	"testing"
 )
 
-// TestEngineAWFEnableApiProxy tests that Claude engine includes --enable-api-proxy
-// in AWF commands, while Copilot does not.
+// TestEngineAWFEnableApiProxy tests that engines with supportsLLMGateway: true
+// include --enable-api-proxy in AWF commands, while engines with supportsLLMGateway: false do not.
 func TestEngineAWFEnableApiProxy(t *testing.T) {
-	t.Run("Claude AWF command includes enable-api-proxy flag", func(t *testing.T) {
+	t.Run("Claude AWF command includes enable-api-proxy flag (supportsLLMGateway: true)", func(t *testing.T) {
 		workflowData := &WorkflowData{
 			Name: "test-workflow",
 			EngineConfig: &EngineConfig{
@@ -31,11 +31,11 @@ func TestEngineAWFEnableApiProxy(t *testing.T) {
 		stepContent := strings.Join(steps[0], "\n")
 
 		if !strings.Contains(stepContent, "--enable-api-proxy") {
-			t.Error("Expected Claude AWF command to contain '--enable-api-proxy' flag")
+			t.Error("Expected Claude AWF command to contain '--enable-api-proxy' flag (supportsLLMGateway: true)")
 		}
 	})
 
-	t.Run("Copilot AWF command does not include enable-api-proxy flag", func(t *testing.T) {
+	t.Run("Copilot AWF command does not include enable-api-proxy flag (supportsLLMGateway: false)", func(t *testing.T) {
 		workflowData := &WorkflowData{
 			Name: "test-workflow",
 			EngineConfig: &EngineConfig{
@@ -58,7 +58,34 @@ func TestEngineAWFEnableApiProxy(t *testing.T) {
 		stepContent := strings.Join(steps[0], "\n")
 
 		if strings.Contains(stepContent, "--enable-api-proxy") {
-			t.Error("Expected Copilot AWF command to NOT contain '--enable-api-proxy' flag")
+			t.Error("Expected Copilot AWF command to NOT contain '--enable-api-proxy' flag (supportsLLMGateway: false)")
+		}
+	})
+
+	t.Run("Codex AWF command does not include enable-api-proxy flag (supportsLLMGateway: false)", func(t *testing.T) {
+		workflowData := &WorkflowData{
+			Name: "test-workflow",
+			EngineConfig: &EngineConfig{
+				ID: "codex",
+			},
+			NetworkPermissions: &NetworkPermissions{
+				Firewall: &FirewallConfig{
+					Enabled: true,
+				},
+			},
+		}
+
+		engine := NewCodexEngine()
+		steps := engine.GetExecutionSteps(workflowData, "test.log")
+
+		if len(steps) == 0 {
+			t.Fatal("Expected at least one execution step")
+		}
+
+		stepContent := strings.Join(steps[0], "\n")
+
+		if strings.Contains(stepContent, "--enable-api-proxy") {
+			t.Error("Expected Codex AWF command to NOT contain '--enable-api-proxy' flag (supportsLLMGateway: false)")
 		}
 	})
 }