fix(steps): expose shell and atmos outputs to later steps

cmd/cmd_utils.go

@@ -33,6 +33,7 @@ import (
 	"github.com/cloudposse/atmos/pkg/dependencies"
 	envpkg "github.com/cloudposse/atmos/pkg/env"
 	pkgFlags "github.com/cloudposse/atmos/pkg/flags"
+	ioLayer "github.com/cloudposse/atmos/pkg/io"
 	l "github.com/cloudposse/atmos/pkg/list"
 	log "github.com/cloudposse/atmos/pkg/logger"
 	"github.com/cloudposse/atmos/pkg/perf"
@@ -1098,32 +1099,43 @@ func executeCustomCommand(
 		// The dispatch is wrapped in a collapsible CI log group (no-op outside
 		// CI / when disabled), labeled with the step name or command. Exec steps
 		// run bare because a successful Unix exec never returns to close a group.
+		var commandResult *stepPkg.StepResult
+		runCommandStep := func(run func(stdout, stderr io.Writer) error) error {
+			var runErr error
+			commandResult, runErr = stepPkg.ExecuteCommandResult(step.Name, run)
+			return runErr
+		}
 		runStep := func() error {
 			switch stepType {
 			case "shell":
 				// Execute shell command (backward compatible).
 				// Steps with tty/interactive attach the user's terminal so commands
 				// like `aws ssm start-session` get a real TTY and own Ctrl-C.
 				commandName := fmt.Sprintf("%s-step-%d", commandConfig.Name, i)
-				return process.RunShellStep(context.Background(), &process.ShellSessionSpec{
-					Command:     commandToRun,
-					Name:        commandName,
-					Dir:         stepWorkDir,
-					Env:         env,
-					TTY:         step.Tty,
-					Interactive: step.Interactive,
-				}, func() error {
-					if step.Output == string(stepPkg.OutputModeNone) {
+				return runCommandStep(func(stdoutCapture, stderrCapture io.Writer) error {
+					return process.RunShellStep(context.Background(), &process.ShellSessionSpec{
+						Command:     commandToRun,
+						Name:        commandName,
+						Dir:         stepWorkDir,
+						Env:         env,
+						TTY:         step.Tty,
+						Interactive: step.Interactive,
+					}, func() error {
+						stdout := ioLayer.MaskWriter(os.Stdout)
+						stderr := ioLayer.MaskWriter(os.Stderr)
+						if step.Output == string(stepPkg.OutputModeNone) {
+							stdout = io.Discard
+							stderr = io.Discard
+						}
 						return e.ExecuteShellWithWriters(&e.ExecuteShellSpec{
 							Command: commandToRun,
 							Name:    commandName,
 							Dir:     stepWorkDir,
 							EnvVars: env,
-							Stdout:  io.Discard,
-							Stderr:  io.Discard,
+							Stdout:  io.MultiWriter(stdout, stdoutCapture),
+							Stderr:  io.MultiWriter(stderr, stderrCapture),
 						})
-					}
-					return e.ExecuteShell(commandToRun, commandName, stepWorkDir, env, false)
+					})
 				})
 			case schema.TaskTypeExec:
 				// Replace the Atmos process with the command (shell exec semantics).
@@ -1140,7 +1152,29 @@ func executeCustomCommand(
 				if execErr != nil {
 					return execErr
 				}
-				return e.ExecuteShellCommand(atmosConfig, execPath, args, stepWorkDir, env, false, "")
+				return runCommandStep(func(stdout, stderr io.Writer) error {
+					execOpts := []e.ShellCommandOption{
+						e.WithStdoutCapture(stdout),
+						e.WithStderrCapture(stderr),
+					}
+					if step.Output == string(stepPkg.OutputModeNone) {
+						execOpts = append(execOpts, e.WithProcessStreams(process.Streams{
+							Stdin:  os.Stdin,
+							Stdout: io.Discard,
+							Stderr: io.Discard,
+						}))
+					}
+					return e.ExecuteShellCommand(
+						atmosConfig,
+						execPath,
+						args,
+						stepWorkDir,
+						env,
+						false,
+						"",
+						execOpts...,
+					)
+				})
 			default:
 				// Check if this is an extended step type (input, confirm, choose, etc.).
 				if stepPkg.IsExtendedStepType(stepType) {
@@ -1175,9 +1209,13 @@ func executeCustomCommand(
 		}
 		err = stepPkg.RunGroupedForType(&atmosConfig, step.Name, commandToRun, stepType, func() error {
 			if step.Retry != nil {
-				return retry.Do(context.Background(), step.Retry, runStep)
+				if retryErr := retry.Do(context.Background(), step.Retry, runStep); retryErr != nil {
+					return retryErr
+				}
+			} else if runErr := runStep(); runErr != nil {
+				return runErr
 			}
-			return runStep()
+			return stepPkg.StoreCommandResult(stepVars, step.Name, step.Outputs, commandResult)
 		})
 		if err != nil {
 			var silentExit errUtils.ExitCodeError

cmd/custom_command_integration_test.go

@@ -429,6 +429,9 @@ func TestCustomCommandShellOutputNoneSuppressesOutput(t *testing.T) {
 
 	tmpDir := t.TempDir()
 	outputFile := filepath.Join(tmpDir, "ran.txt")
+	resultFile := filepath.Join(tmpDir, "result.env")
+	exePath, err := os.Executable()
+	require.NoError(t, err)
 
 	atmosConfig := schema.AtmosConfiguration{
 		BasePath: tmpDir,
@@ -441,14 +444,23 @@ func TestCustomCommandShellOutputNoneSuppressesOutput(t *testing.T) {
 						Type:    "shell",
 						Name:    "quiet",
 						Output:  "none",
-						Command: fmt.Sprintf("echo stdout-visible; echo stderr-visible >&2; printf ran > %q", outputFile),
+						Command: fmt.Sprintf("printf stdout-visible; printf stderr-visible >&2; printf ran > %q", outputFile),
+					},
+					{
+						Type:    "shell",
+						Output:  "none",
+						Command: fmt.Sprintf("%q", exePath),
+						Env: map[string]string{
+							"_ATMOS_TEST_DUMP_ENV": resultFile,
+							"CAPTURED_RESULT":      "{{ .steps.quiet.value }}|{{ .steps.quiet.metadata.stdout }}|{{ .steps.quiet.metadata.stderr }}|{{ .steps.quiet.metadata.exit_code }}",
+						},
 					},
 				},
 			},
 		},
 	}
 
-	err := processCustomCommands(atmosConfig, atmosConfig.Commands, RootCmd)
+	err = processCustomCommands(atmosConfig, atmosConfig.Commands, RootCmd)
 	require.NoError(t, err)
 
 	customCmd, _, err := RootCmd.Find([]string{"test-output-none"})
@@ -465,6 +477,9 @@ func TestCustomCommandShellOutputNoneSuppressesOutput(t *testing.T) {
 	actual, err := os.ReadFile(outputFile)
 	require.NoError(t, err)
 	assert.Equal(t, "ran", string(actual))
+	resultEnv, err := os.ReadFile(resultFile)
+	require.NoError(t, err)
+	assert.Equal(t, "stdout-visible|stdout-visible|stderr-visible|0", extractEnvVar(string(resultEnv), "CAPTURED_RESULT"))
 }
 
 // TestCustomCommandIntegration_MockProviderEnvironment tests that custom commands with mock provider
@@ -761,6 +776,9 @@ func TestCustomCommandIntegration_RetriesShellStep(t *testing.T) {
 
 	tmpDir := t.TempDir()
 	attemptsFile := filepath.Join(tmpDir, "attempts.txt")
+	resultFile := filepath.Join(tmpDir, "result.env")
+	exePath, err := os.Executable()
+	require.NoError(t, err)
 	maxAttempts := 2
 	initialDelay := time.Millisecond
 
@@ -769,6 +787,7 @@ func TestCustomCommandIntegration_RetriesShellStep(t *testing.T) {
 		Description: "Test retry shell step",
 		Steps: schema.Tasks{
 			{
+				Name:    "retry",
 				Command: customCommandRetryHelperCommand(t, attemptsFile),
 				Type:    "shell",
 				Retry: &schema.RetryConfig{
@@ -777,6 +796,15 @@ func TestCustomCommandIntegration_RetriesShellStep(t *testing.T) {
 					BackoffStrategy: "constant",
 				},
 			},
+			{
+				Command: fmt.Sprintf("%q", exePath),
+				Type:    "shell",
+				Output:  "none",
+				Env: map[string]string{
+					"_ATMOS_TEST_DUMP_ENV": resultFile,
+					"CAPTURED_RESULT":      "{{ .steps.retry.value }}|{{ .steps.retry.metadata.stdout }}|{{ .steps.retry.metadata.stderr }}",
+				},
+			},
 		},
 	}
 	atmosConfig.Commands = []schema.Command{testCommand}
@@ -798,6 +826,9 @@ func TestCustomCommandIntegration_RetriesShellStep(t *testing.T) {
 	attempts, err := os.ReadFile(attemptsFile)
 	require.NoError(t, err)
 	assert.Equal(t, "2", strings.TrimSpace(string(attempts)))
+	resultEnv, err := os.ReadFile(resultFile)
+	require.NoError(t, err)
+	assert.Equal(t, "attempt-2|attempt-2|warning-2", extractEnvVar(string(resultEnv), "CAPTURED_RESULT"))
 }
 
 func TestCustomCommandIntegration_ShellStepWithoutRetryRunsOnce(t *testing.T) {
@@ -1051,6 +1082,11 @@ func customCommandAttemptHelperCommand(t *testing.T, path string) string {
 	return fmt.Sprintf("%q -test.run=TestCustomCommandIntegrationAttemptHelper -- %s", exe, encodedPath)
 }
 
+func customCommandAtmosAttemptHelperArgs(path string) string {
+	encodedPath := base64.RawURLEncoding.EncodeToString([]byte(path))
+	return fmt.Sprintf("-test.run=TestCustomCommandIntegrationAttemptHelper -- %s", encodedPath)
+}
+
 func TestCustomCommandIntegrationWriteHelper(t *testing.T) {
 	separator := -1
 	for i, arg := range os.Args {
@@ -1098,6 +1134,8 @@ func TestCustomCommandIntegrationRetryHelper(t *testing.T) {
 		attempt = parsed + 1
 	}
 	require.NoError(t, os.WriteFile(path, []byte(strconv.Itoa(attempt)), 0o600))
+	_, _ = fmt.Fprintf(os.Stdout, "attempt-%d", attempt)
+	_, _ = fmt.Fprintf(os.Stderr, "warning-%d", attempt)
 	if attempt < 2 {
 		os.Exit(1)
 	}