cloudposse
diff --git a/‎cmd/ai/flags/flags.go‎
Lines changed: 130 additions & 0 deletions b/‎cmd/ai/flags/flags.go‎
Lines changed: 130 additions & 0 deletions
diff --git a/‎cmd/ai/flags/flags_test.go‎
Lines changed: 194 additions & 0 deletions b/‎cmd/ai/flags/flags_test.go‎
Lines changed: 194 additions & 0 deletions
diff --git a/‎cmd/ai/setup/setup.go‎
Lines changed: 55 additions & 0 deletions b/‎cmd/ai/setup/setup.go‎
Lines changed: 55 additions & 0 deletions
@@ -0,0 +1,130 @@
+package flags
+
+import (
+	"os"
+	"strconv"
+	"strings"
+
+	"github.com/cloudposse/atmos/pkg/perf"
+)
+
+// HasAIFlag checks if --ai flag is present in os.Args.
+// This is needed before Cobra parses flags because we set up output capture
+// in Execute() before internal.Execute() runs the command.
+func HasAIFlag() bool {
+	defer perf.Track(nil, "flags.HasAIFlag")()
+
+	return HasAIFlagInternal(os.Args)
+}
+
+// HasAIFlagInternal checks if --ai flag is present in the provided args.
+// Uses last-value-wins semantics to match Cobra's behavior for repeated flags.
+// Only falls back to ATMOS_AI environment variable when no CLI flag is present.
+func HasAIFlagInternal(args []string) bool {
+	var explicit *bool
+	for _, arg := range args {
+		// Stop scanning after bare "--" (end-of-flags delimiter).
+		if arg == "--" {
+			break
+		}
+		// Bare --ai is equivalent to --ai=true.
+		if arg == "--ai" {
+			v := true
+			explicit = &v
+			continue
+		}
+		// Explicit --ai=<value>: parse the boolean to respect --ai=false.
+		if strings.HasPrefix(arg, "--ai=") {
+			val, err := strconv.ParseBool(strings.TrimPrefix(arg, "--ai="))
+			if err != nil {
+				continue
+			}
+			explicit = &val
+		}
+	}
+	if explicit != nil {
+		return *explicit
+	}
+	// Fall back to ATMOS_AI environment variable for CI/CD env-only usage.
+	//nolint:forbidigo // Must use os.Getenv: AI flag is processed before Viper configuration loads.
+	val, err := strconv.ParseBool(os.Getenv("ATMOS_AI"))
+	return err == nil && val
+}
+
+// ParseSkillFlag extracts all --skill flag values from os.Args.
+// This is needed before Cobra parses flags because we validate and load skills
+// in Execute() before internal.Execute() runs the command.
+func ParseSkillFlag() []string {
+	defer perf.Track(nil, "flags.ParseSkillFlag")()
+
+	return ParseSkillFlagInternal(os.Args)
+}
+
+// ParseSkillFlagInternal extracts all --skill flag values from the provided args.
+// Supports repeated flags (--skill a --skill b) and comma-separated values (--skill a,b).
+// Respects CLI precedence: if --skill is explicitly provided (even as empty), the env var is not consulted.
+// Only falls back to ATMOS_SKILL environment variable when no --skill flag is present.
+func ParseSkillFlagInternal(args []string) []string {
+	var result []string
+	flagSeen := false
+	for i, arg := range args {
+		// Stop scanning after bare "--" (end-of-flags delimiter).
+		if arg == "--" {
+			break
+		}
+
+		var value string
+		if arg == "--skill" {
+			flagSeen = true
+			if i+1 < len(args) && !strings.HasPrefix(args[i+1], "-") {
+				value = args[i+1]
+			}
+		} else if strings.HasPrefix(arg, "--skill=") {
+			value = strings.TrimPrefix(arg, "--skill=")
+			flagSeen = true
+		}
+
+		if value != "" {
+			result = append(result, SplitCSV(value)...)
+		}
+	}
+
+	// Fall back to ATMOS_SKILL environment variable only when no --skill CLI flag was provided.
+	if !flagSeen {
+		//nolint:forbidigo // Must use os.Getenv: skill flag is processed before Viper configuration loads.
+		result = SplitCSV(os.Getenv("ATMOS_SKILL"))
+	}
+
+	return result
+}
+
+// HasHelpFlag checks if --help or -h is present in os.Args (before the "--" delimiter).
+// Used to short-circuit AI setup when the user is requesting help.
+func HasHelpFlag() bool {
+	return HasHelpFlagInternal(os.Args)
+}
+
+// HasHelpFlagInternal checks if --help, -h, or the "help" subcommand is present in the provided args.
+func HasHelpFlagInternal(args []string) bool {
+	for _, arg := range args {
+		if arg == "--" {
+			break
+		}
+		if arg == "--help" || arg == "-h" || arg == "help" {
+			return true
+		}
+	}
+	return false
+}
+
+// SplitCSV splits a comma-separated string into trimmed, non-empty values.
+func SplitCSV(value string) []string {
+	var result []string
+	for _, v := range strings.Split(value, ",") {
+		v = strings.TrimSpace(v)
+		if v != "" {
+			result = append(result, v)
+		}
+	}
+	return result
+}
@@ -0,0 +1,194 @@
+package flags
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// TestHasAIFlagInternal tests the HasAIFlagInternal function that parses --ai from os.Args.
+func TestHasAIFlagInternal(t *testing.T) {
+	tests := []struct {
+		name     string
+		args     []string
+		expected bool
+	}{
+		{name: "no --ai flag", args: []string{"atmos", "terraform", "plan"}, expected: false},
+		{name: "--ai flag present", args: []string{"atmos", "--ai", "terraform", "plan"}, expected: true},
+		{name: "--ai flag at end", args: []string{"atmos", "terraform", "plan", "--ai"}, expected: true},
+		{name: "--ai=true", args: []string{"atmos", "--ai=true", "terraform", "plan"}, expected: true},
+		{name: "--ai=false is not enabled", args: []string{"atmos", "--ai=false", "terraform", "plan"}, expected: false},
+		{name: "--ai after -- delimiter is ignored", args: []string{"atmos", "terraform", "plan", "--", "--ai"}, expected: false},
+		{name: "similar flag --aim is not matched", args: []string{"atmos", "--aim", "terraform", "plan"}, expected: false},
+		{name: "empty args", args: []string{}, expected: false},
+		{name: "repeated --ai last wins (true then false)", args: []string{"atmos", "--ai", "--ai=false", "terraform", "plan"}, expected: false},
+		{name: "repeated --ai last wins (false then true)", args: []string{"atmos", "--ai=false", "--ai", "terraform", "plan"}, expected: true},
+		{name: "repeated --ai=value last wins", args: []string{"atmos", "--ai=true", "--ai=false", "terraform"}, expected: false},
+		{name: "--ai=invalid skipped, bare --ai wins", args: []string{"atmos", "--ai=maybe", "--ai", "terraform"}, expected: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := HasAIFlagInternal(tt.args)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+// TestHasAIFlagInternal_EnvVarFallback tests env var fallback behavior for --ai detection.
+func TestHasAIFlagInternal_EnvVarFallback(t *testing.T) {
+	tests := []struct {
+		name     string
+		envValue string
+		expected bool
+	}{
+		{name: "ATMOS_AI=true enables AI", envValue: "true", expected: true},
+		{name: "ATMOS_AI=1 enables AI", envValue: "1", expected: true},
+		{name: "ATMOS_AI=false disables AI", envValue: "false", expected: false},
+		{name: "ATMOS_AI=0 disables AI", envValue: "0", expected: false},
+		{name: "ATMOS_AI=invalid disables AI", envValue: "invalid", expected: false},
+		{name: "ATMOS_AI empty disables AI", envValue: "", expected: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Setenv("ATMOS_AI", tt.envValue)
+			result := HasAIFlagInternal([]string{"atmos", "terraform", "plan"})
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+// TestHasAIFlagInternal_CLIOverridesEnv verifies CLI flag takes precedence over env var.
+func TestHasAIFlagInternal_CLIOverridesEnv(t *testing.T) {
+	t.Setenv("ATMOS_AI", "true")
+	result := HasAIFlagInternal([]string{"atmos", "--ai=false", "terraform", "plan"})
+	assert.False(t, result, "--ai=false should override ATMOS_AI=true")
+}
+
+// TestHasAIFlagInternal_InvalidBoolValue tests --ai with invalid boolean value.
+func TestHasAIFlagInternal_InvalidBoolValue(t *testing.T) {
+	result := HasAIFlagInternal([]string{"atmos", "--ai=maybe", "terraform", "plan"})
+	assert.False(t, result, "--ai=maybe should return false (invalid boolean)")
+}
+
+// TestParseSkillFlagInternal tests the ParseSkillFlagInternal function.
+func TestParseSkillFlagInternal(t *testing.T) {
+	tests := []struct {
+		name     string
+		args     []string
+		expected []string
+	}{
+		{name: "no --skill flag", args: []string{"atmos", "terraform", "plan"}, expected: nil},
+		{name: "--skill with separate value", args: []string{"atmos", "--ai", "--skill", "atmos-terraform", "terraform", "plan"}, expected: []string{"atmos-terraform"}},
+		{name: "--skill=value format", args: []string{"atmos", "--ai", "--skill=atmos-stacks", "describe", "stacks"}, expected: []string{"atmos-stacks"}},
+		{name: "--skill at end with value", args: []string{"atmos", "terraform", "plan", "--ai", "--skill", "atmos-terraform"}, expected: []string{"atmos-terraform"}},
+		{name: "--skill after -- delimiter is ignored", args: []string{"atmos", "terraform", "plan", "--", "--skill", "atmos-terraform"}, expected: nil},
+		{name: "--skill without value (next arg is flag)", args: []string{"atmos", "--skill", "--ai", "terraform", "plan"}, expected: nil},
+		{name: "--skill at end without value", args: []string{"atmos", "terraform", "plan", "--skill"}, expected: nil},
+		{name: "empty args", args: []string{}, expected: nil},
+		{name: "--skill=empty value", args: []string{"atmos", "--skill=", "terraform", "plan"}, expected: nil},
+		{name: "--skilled is not matched", args: []string{"atmos", "--skilled", "atmos-terraform", "terraform", "plan"}, expected: nil},
+		{name: "--skilled=value not matched", args: []string{"atmos", "--skilled=atmos-terraform", "terraform", "plan"}, expected: nil},
+		{name: "--skill between other flags", args: []string{"atmos", "--logs-level=Debug", "--skill", "atmos-validation", "--ai", "validate", "stacks"}, expected: []string{"atmos-validation"}},
+		{name: "hyphens in skill name", args: []string{"atmos", "--skill=my-custom-skill-v2", "terraform", "plan"}, expected: []string{"my-custom-skill-v2"}},
+		{name: "multiple --skill flags", args: []string{"atmos", "--skill", "first", "--skill", "second", "terraform", "plan"}, expected: []string{"first", "second"}},
+		{name: "comma-separated skills", args: []string{"atmos", "--ai", "--skill", "atmos-terraform,atmos-stacks", "terraform", "plan"}, expected: []string{"atmos-terraform", "atmos-stacks"}},
+		{name: "comma-separated with =", args: []string{"atmos", "--ai", "--skill=atmos-terraform,atmos-stacks", "terraform", "plan"}, expected: []string{"atmos-terraform", "atmos-stacks"}},
+		{name: "mixed repeated and comma", args: []string{"atmos", "--skill", "a,b", "--skill", "c", "terraform", "plan"}, expected: []string{"a", "b", "c"}},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := ParseSkillFlagInternal(tt.args)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+// TestParseSkillFlagInternal_EnvVarFallback tests env var fallback for --skill.
+func TestParseSkillFlagInternal_EnvVarFallback(t *testing.T) {
+	tests := []struct {
+		name     string
+		envValue string
+		expected []string
+	}{
+		{name: "single skill from env", envValue: "atmos-terraform", expected: []string{"atmos-terraform"}},
+		{name: "comma-separated from env", envValue: "a,b,c", expected: []string{"a", "b", "c"}},
+		{name: "empty env", envValue: "", expected: nil},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Setenv("ATMOS_SKILL", tt.envValue)
+			result := ParseSkillFlagInternal([]string{"atmos", "terraform", "plan"})
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+// TestParseSkillFlagInternal_CLIOverridesEnv verifies CLI flag prevents env var fallback.
+func TestParseSkillFlagInternal_CLIOverridesEnv(t *testing.T) {
+	t.Setenv("ATMOS_SKILL", "from-env")
+	result := ParseSkillFlagInternal([]string{"atmos", "--skill", "from-cli", "terraform", "plan"})
+	assert.Equal(t, []string{"from-cli"}, result, "CLI --skill should override ATMOS_SKILL env var")
+}
+
+// TestParseSkillFlagInternal_BareSkillSuppressesEnv verifies bare --skill without value suppresses env fallback.
+func TestParseSkillFlagInternal_BareSkillSuppressesEnv(t *testing.T) {
+	t.Setenv("ATMOS_SKILL", "from-env")
+	result := ParseSkillFlagInternal([]string{"atmos", "--skill", "--ai", "terraform", "plan"})
+	assert.Nil(t, result, "bare --skill without value should suppress ATMOS_SKILL env var fallback")
+}
+
+// TestHasHelpFlagInternal tests the HasHelpFlagInternal function.
+func TestHasHelpFlagInternal(t *testing.T) {
+	tests := []struct {
+		name     string
+		args     []string
+		expected bool
+	}{
+		{name: "no help flag", args: []string{"atmos", "terraform", "plan"}, expected: false},
+		{name: "--help flag", args: []string{"atmos", "terraform", "plan", "--help"}, expected: true},
+		{name: "-h flag", args: []string{"atmos", "terraform", "-h"}, expected: true},
+		{name: "help subcommand", args: []string{"atmos", "help", "terraform"}, expected: true},
+		{name: "--help after -- is ignored", args: []string{"atmos", "terraform", "--", "--help"}, expected: false},
+		{name: "-h after -- is ignored", args: []string{"atmos", "terraform", "--", "-h"}, expected: false},
+		{name: "empty args", args: []string{}, expected: false},
+		{name: "--help with --ai", args: []string{"atmos", "--ai", "terraform", "plan", "--help"}, expected: true},
+		{name: "--help with --skill", args: []string{"atmos", "--skill", "atmos-terraform", "--help"}, expected: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := HasHelpFlagInternal(tt.args)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+// TestSplitCSV tests the SplitCSV helper function.
+func TestSplitCSV(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected []string
+	}{
+		{name: "empty string", input: "", expected: nil},
+		{name: "single value", input: "a", expected: []string{"a"}},
+		{name: "two values", input: "a,b", expected: []string{"a", "b"}},
+		{name: "values with whitespace", input: " a , b , c ", expected: []string{"a", "b", "c"}},
+		{name: "trailing comma", input: "a,b,", expected: []string{"a", "b"}},
+		{name: "leading comma", input: ",a,b", expected: []string{"a", "b"}},
+		{name: "multiple commas", input: "a,,b", expected: []string{"a", "b"}},
+		{name: "only commas", input: ",,", expected: nil},
+		{name: "whitespace only values", input: " , , ", expected: nil},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := SplitCSV(tt.input)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
@@ -0,0 +1,55 @@
+package setup
+
+import (
+	"strings"
+
+	"github.com/spf13/cobra"
+
+	aiflags "github.com/cloudposse/atmos/cmd/ai/flags"
+	errUtils "github.com/cloudposse/atmos/errors"
+	"github.com/cloudposse/atmos/pkg/ai/analyze"
+	"github.com/cloudposse/atmos/pkg/perf"
+	"github.com/cloudposse/atmos/pkg/schema"
+)
+
+// InitAI parses AI flags from os.Args, validates configuration, and starts output capture.
+// Returns a ready-to-use Context or an error if validation fails.
+// The caller MUST call ctx.Cleanup() (via defer) to restore stdout/stderr.
+func InitAI(atmosConfig *schema.AtmosConfiguration) (*analyze.Context, error) {
+	defer perf.Track(nil, "setup.InitAI")()
+
+	// Short-circuit for help requests: --help/-h/help should never fail due to AI env vars.
+	if aiflags.HasHelpFlag() {
+		return analyze.NewDisabledContext(), nil
+	}
+
+	aiEnabled := aiflags.HasAIFlag()
+	skillNames := aiflags.ParseSkillFlag()
+
+	// Validate --skill requires --ai.
+	if len(skillNames) > 0 && !aiEnabled {
+		skillList := strings.Join(skillNames, ",")
+		return nil, errUtils.Build(errUtils.ErrAISkillRequiresAIFlag).
+			WithExplanation("The --skill flag provides domain-specific context for AI analysis, but AI analysis is not enabled. Use --skill together with --ai.").
+			WithHintf("Add --ai to enable AI analysis:\n  atmos <command> --ai --skill %s", skillList).
+			WithHintf("Or use environment variables:\n  ATMOS_AI=true ATMOS_SKILL=%s atmos <command>", skillList).
+			Err()
+	}
+
+	if !aiEnabled {
+		return analyze.NewDisabledContext(), nil
+	}
+
+	return analyze.Setup(atmosConfig, skillNames, analyze.BuildCommandName())
+}
+
+// IsAISubcommand checks whether cmd is the "atmos ai" subcommand (or one of its children).
+// Used after Cobra execution to skip AI analysis for AI commands (avoid double processing).
+func IsAISubcommand(cmd *cobra.Command) bool {
+	for c := cmd; c != nil; c = c.Parent() {
+		if c.Name() == "ai" && c.Parent() != nil && c.Parent().Parent() == nil {
+			return true
+		}
+	}
+	return false
+}