Extract duplicate custom engine step handling into shared helper functions

pkg/workflow/claude_engine.go

@@ -86,19 +86,8 @@ func (e *ClaudeEngine) GetVersionCommand() string {
 
 // GetExecutionSteps returns the GitHub Actions steps for executing Claude
 func (e *ClaudeEngine) GetExecutionSteps(workflowData *WorkflowData, logFile string) []GitHubActionStep {
-	var steps []GitHubActionStep
-
 	// Handle custom steps if they exist in engine config
-	if workflowData.EngineConfig != nil && len(workflowData.EngineConfig.Steps) > 0 {
-		for _, step := range workflowData.EngineConfig.Steps {
-			stepYAML, err := e.convertStepToYAML(step)
-			if err != nil {
-				// Log error but continue with other steps
-				continue
-			}
-			steps = append(steps, GitHubActionStep{stepYAML})
-		}
-	}
+	steps := InjectCustomEngineSteps(workflowData, e.convertStepToYAML)
 
 	// Build claude CLI arguments based on configuration
 	var claudeArgs []string
@@ -616,14 +605,8 @@ func (e *ClaudeEngine) RenderMCPConfig(yaml *strings.Builder, tools map[string]a
 		case "web-fetch":
 			renderMCPFetchServerConfig(yaml, "json", "              ", isLast, false)
 		default:
-			// Handle custom MCP tools (those with MCP-compatible type)
-			if toolConfig, ok := tools[toolName].(map[string]any); ok {
-				if hasMcp, _ := hasMCPConfig(toolConfig); hasMcp {
-					if err := e.renderClaudeMCPConfig(yaml, toolName, toolConfig, isLast); err != nil {
-						fmt.Printf("Error generating custom MCP configuration for %s: %v\n", toolName, err)
-					}
-				}
-			}
+			// Handle custom MCP tools using shared helper
+			HandleCustomMCPToolInSwitch(yaml, toolName, tools, isLast, e.renderClaudeMCPConfig)
 		}
 	}
 

pkg/workflow/codex_engine.go

@@ -95,19 +95,8 @@ func (e *CodexEngine) GetDeclaredOutputFiles() []string {
 
 // GetExecutionSteps returns the GitHub Actions steps for executing Codex
 func (e *CodexEngine) GetExecutionSteps(workflowData *WorkflowData, logFile string) []GitHubActionStep {
-	var steps []GitHubActionStep
-
 	// Handle custom steps if they exist in engine config
-	if workflowData.EngineConfig != nil && len(workflowData.EngineConfig.Steps) > 0 {
-		for _, step := range workflowData.EngineConfig.Steps {
-			stepYAML, err := e.convertStepToYAML(step)
-			if err != nil {
-				// Log error but continue with other steps
-				continue
-			}
-			steps = append(steps, GitHubActionStep{stepYAML})
-		}
-	}
+	steps := InjectCustomEngineSteps(workflowData, e.convertStepToYAML)
 
 	// Build model parameter only if specified in engineConfig
 	var modelParam string
@@ -253,14 +242,10 @@ func (e *CodexEngine) RenderMCPConfig(yaml *strings.Builder, tools map[string]an
 		case "web-fetch":
 			renderMCPFetchServerConfig(yaml, "toml", "          ", false, false)
 		default:
-			// Handle custom MCP tools (those with MCP-compatible type)
-			if toolConfig, ok := expandedTools[toolName].(map[string]any); ok {
-				if hasMcp, _ := hasMCPConfig(toolConfig); hasMcp {
-					if err := e.renderCodexMCPConfig(yaml, toolName, toolConfig); err != nil {
-						fmt.Printf("Error generating custom MCP configuration for %s: %v\n", toolName, err)
-					}
-				}
-			}
+			// Handle custom MCP tools using shared helper (with adapter for isLast parameter)
+			HandleCustomMCPToolInSwitch(yaml, toolName, expandedTools, false, func(yaml *strings.Builder, toolName string, toolConfig map[string]any, isLast bool) error {
+				return e.renderCodexMCPConfig(yaml, toolName, toolConfig)
+			})
 		}
 	}
 

pkg/workflow/copilot_engine.go

@@ -65,19 +65,8 @@ func (e *CopilotEngine) GetVersionCommand() string {
 
 // GetExecutionSteps returns the GitHub Actions steps for executing GitHub Copilot CLI
 func (e *CopilotEngine) GetExecutionSteps(workflowData *WorkflowData, logFile string) []GitHubActionStep {
-	var steps []GitHubActionStep
-
 	// Handle custom steps if they exist in engine config
-	if workflowData.EngineConfig != nil && len(workflowData.EngineConfig.Steps) > 0 {
-		for _, step := range workflowData.EngineConfig.Steps {
-			stepYAML, err := e.convertStepToYAML(step)
-			if err != nil {
-				// Log error but continue with other steps
-				continue
-			}
-			steps = append(steps, GitHubActionStep{stepYAML})
-		}
-	}
+	steps := InjectCustomEngineSteps(workflowData, e.convertStepToYAML)
 
 	// Build copilot CLI arguments based on configuration
 	var copilotArgs = []string{"--add-dir", "/tmp/gh-aw/", "--log-level", "all", "--log-dir", logsFolder}
@@ -252,14 +241,8 @@ func (e *CopilotEngine) RenderMCPConfig(yaml *strings.Builder, tools map[string]
 		case "web-fetch":
 			renderMCPFetchServerConfig(yaml, "json", "              ", isLast, true)
 		default:
-			// Handle custom MCP tools (those with MCP-compatible type)
-			if toolConfig, ok := tools[toolName].(map[string]any); ok {
-				if hasMcp, _ := hasMCPConfig(toolConfig); hasMcp {
-					if err := e.renderCopilotMCPConfig(yaml, toolName, toolConfig, isLast); err != nil {
-						fmt.Printf("Error generating custom MCP configuration for %s: %v\n", toolName, err)
-					}
-				}
-			}
+			// Handle custom MCP tools using shared helper
+			HandleCustomMCPToolInSwitch(yaml, toolName, tools, isLast, e.renderCopilotMCPConfig)
 		}
 	}
 

pkg/workflow/custom_engine.go

@@ -172,14 +172,8 @@ func (e *CustomEngine) RenderMCPConfig(yaml *strings.Builder, tools map[string]a
 		case "web-fetch":
 			renderMCPFetchServerConfig(yaml, "json", "              ", isLast, false)
 		default:
-			// Handle custom MCP tools (those with MCP-compatible type)
-			if toolConfig, ok := tools[toolName].(map[string]any); ok {
-				if hasMcp, _ := hasMCPConfig(toolConfig); hasMcp {
-					if err := e.renderCustomMCPConfig(yaml, toolName, toolConfig, isLast); err != nil {
-						fmt.Printf("Error generating custom MCP configuration for %s: %v\n", toolName, err)
-					}
-				}
-			}
+			// Handle custom MCP tools using shared helper
+			HandleCustomMCPToolInSwitch(yaml, toolName, tools, isLast, e.renderCustomMCPConfig)
 		}
 	}
 

pkg/workflow/engine_shared_helpers.go

@@ -0,0 +1,70 @@
+package workflow
+
+import (
+	"fmt"
+	"strings"
+)
+
+// InjectCustomEngineSteps processes custom steps from engine config and converts them to GitHubActionSteps.
+// This shared function extracts the common pattern used by Copilot, Codex, and Claude engines.
+//
+// Parameters:
+//   - workflowData: The workflow data containing engine configuration
+//   - convertStepFunc: A function that converts a step map to YAML string (engine-specific)
+//
+// Returns:
+//   - []GitHubActionStep: Array of custom steps ready to be included in the execution pipeline
+func InjectCustomEngineSteps(
+	workflowData *WorkflowData,
+	convertStepFunc func(map[string]any) (string, error),
+) []GitHubActionStep {
+	var steps []GitHubActionStep
+
+	// Handle custom steps if they exist in engine config
+	if workflowData.EngineConfig != nil && len(workflowData.EngineConfig.Steps) > 0 {
+		for _, step := range workflowData.EngineConfig.Steps {
+			stepYAML, err := convertStepFunc(step)
+			if err != nil {
+				// Log error but continue with other steps
+				continue
+			}
+			steps = append(steps, GitHubActionStep{stepYAML})
+		}
+	}
+
+	return steps
+}
+
+// RenderCustomMCPToolConfigHandler is a function type that engines must provide to render their specific MCP config
+type RenderCustomMCPToolConfigHandler func(yaml *strings.Builder, toolName string, toolConfig map[string]any, isLast bool) error
+
+// HandleCustomMCPToolInSwitch processes custom MCP tools in the default case of a switch statement.
+// This shared function extracts the common pattern used across all workflow engines.
+//
+// Parameters:
+//   - yaml: The string builder for YAML output
+//   - toolName: The name of the tool being processed
+//   - tools: The tools map containing tool configurations (supports both expanded and non-expanded tools)
+//   - isLast: Whether this is the last tool in the list
+//   - renderFunc: Engine-specific function to render the MCP configuration
+//
+// Returns:
+//   - bool: true if a custom MCP tool was handled, false otherwise
+func HandleCustomMCPToolInSwitch(
+	yaml *strings.Builder,
+	toolName string,
+	tools map[string]any,
+	isLast bool,
+	renderFunc RenderCustomMCPToolConfigHandler,
+) bool {
+	// Handle custom MCP tools (those with MCP-compatible type)
+	if toolConfig, ok := tools[toolName].(map[string]any); ok {
+		if hasMcp, _ := hasMCPConfig(toolConfig); hasMcp {
+			if err := renderFunc(yaml, toolName, toolConfig, isLast); err != nil {
+				fmt.Printf("Error generating custom MCP configuration for %s: %v\n", toolName, err)
+			}
+			return true
+		}
+	}
+	return false
+}