Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 46 additions & 14 deletions cli/azd/cmd/container.go
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@ import (
"github.com/azure/azure-dev/cli/azd/pkg/workflow"
"github.com/mattn/go-colorable"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
)

// Registers a transient action initializer for the specified action name
Expand Down Expand Up @@ -891,13 +892,13 @@ func registerCommonDependencies(container *ioc.NestedContainer) {
container.MustRegisterNamedSingleton(platformName, constructor)
}

container.MustRegisterSingleton(func(s ioc.ServiceLocator) (workflow.AzdCommandRunner, error) {
var rootCmd *cobra.Command
if err := s.ResolveNamed("root-cmd", &rootCmd); err != nil {
return nil, err
container.MustRegisterSingleton(func() workflow.AzdCommandRunner {
return &workflowCmdAdapter{
newCommand: func() *cobra.Command {
return NewRootCmd(false, nil, container)
},
globalArgs: extractGlobalArgs(),
}
return &workflowCmdAdapter{cmd: rootCmd}, nil

})
container.MustRegisterSingleton(workflow.NewRunner)

Expand Down Expand Up @@ -948,19 +949,50 @@ func registerCommonDependencies(container *ioc.NestedContainer) {
registerAction[*configShowAction](container, "azd-config-show-action")
}

// workflowCmdAdapter adapts a cobra command to the workflow.AzdCommandRunner interface
// workflowCmdAdapter adapts a cobra command to the workflow.AzdCommandRunner interface.
// On each ExecuteContext call, a fresh cobra command tree is built via the newCommand factory
// to avoid stale context or command state from previous executions.
// See: https://github.com/Azure/azure-dev/issues/6530
type workflowCmdAdapter struct {
cmd *cobra.Command
newCommand func() *cobra.Command
globalArgs []string
}

func (w *workflowCmdAdapter) SetArgs(args []string) {
w.cmd.SetArgs(args)
// ExecuteContext implements workflow.AzdCommandRunner.
// It rebuilds the cobra command tree on each call to ensure a clean slate,
// preventing "context canceled" errors from stale command state during retries.
// Global flags from the original process invocation are appended to the step args
// so that persistent flags (e.g., --trace-log-file) are properly parsed and visible
// to telemetry middleware on the fresh command tree.
func (w *workflowCmdAdapter) ExecuteContext(ctx context.Context, args []string) error {
childCtx := middleware.WithChildAction(ctx)
rootCmd := w.newCommand()
// Always set args explicitly to prevent Cobra from falling back to os.Args[1:].
// Cobra uses os.Args when cmd.args is nil (but not when it's an empty slice).
mergedArgs := append(slices.Clone(args), w.globalArgs...)
if mergedArgs == nil {
mergedArgs = []string{}
}
rootCmd.SetArgs(mergedArgs)
return rootCmd.ExecuteContext(childCtx)
}

// ExecuteContext implements workflow.AzdCommandRunner
func (w *workflowCmdAdapter) ExecuteContext(ctx context.Context) error {
childCtx := middleware.WithChildAction(ctx)
return w.cmd.ExecuteContext(childCtx)
// extractGlobalArgs extracts global flag arguments from the process command line.
// It parses os.Args against the global flag set and returns only the flags that were
// explicitly set by the user, formatted as command-line arguments.
func extractGlobalArgs() []string {
globalFlagSet := CreateGlobalFlagSet()
globalFlagSet.SetOutput(io.Discard)
globalFlagSet.ParseErrorsAllowlist = pflag.ParseErrorsAllowlist{UnknownFlags: true}
_ = globalFlagSet.Parse(os.Args[1:])

var result []string
globalFlagSet.VisitAll(func(f *pflag.Flag) {
if f.Changed {
result = append(result, fmt.Sprintf("--%s", f.Name), f.Value.String())
}
})
return result
}

// ArmClientInitializer is a function definition for all Azure SDK ARM Client
Expand Down
168 changes: 126 additions & 42 deletions cli/azd/cmd/container_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"testing"

"github.com/azure/azure-dev/cli/azd/cmd/middleware"
"github.com/azure/azure-dev/cli/azd/internal"
"github.com/azure/azure-dev/cli/azd/pkg/environment/azdcontext"
"github.com/azure/azure-dev/cli/azd/pkg/ioc"
"github.com/azure/azure-dev/cli/azd/pkg/lazy"
Expand Down Expand Up @@ -170,30 +171,32 @@ func Test_WorkflowCmdAdapter_ContextPropagation(t *testing.T) {
// Track which contexts were seen by the subcommand
var receivedContexts []context.Context

// Create a root command with a subcommand
rootCmd := &cobra.Command{
Use: "root",
}
// Create a command factory that builds a fresh tree on each call
newCommand := func() *cobra.Command {
rootCmd := &cobra.Command{
Use: "root",
}

subCmd := &cobra.Command{
Use: "sub",
RunE: func(cmd *cobra.Command, args []string) error {
// Capture the context that the subcommand receives
receivedContexts = append(receivedContexts, cmd.Context())
return nil
},
}
subCmd := &cobra.Command{
Use: "sub",
RunE: func(cmd *cobra.Command, args []string) error {
// Capture the context that the subcommand receives
receivedContexts = append(receivedContexts, cmd.Context())
return nil
},
}

rootCmd.AddCommand(subCmd)
rootCmd.AddCommand(subCmd)
return rootCmd
}

// Create the adapter
adapter := &workflowCmdAdapter{cmd: rootCmd}
// Create the adapter with a factory
adapter := &workflowCmdAdapter{newCommand: newCommand}

// In production, main.go wraps with context.WithoutCancel.
// Simulate this by using a non-cancellable context.
ctx := context.WithoutCancel(context.Background())
adapter.SetArgs([]string{"sub"})
err := adapter.ExecuteContext(ctx)
err := adapter.ExecuteContext(ctx, []string{"sub"})
require.NoError(t, err)
require.Len(t, receivedContexts, 1, "Execution should have received context")

Expand All @@ -209,9 +212,8 @@ func Test_WorkflowCmdAdapter_ContextPropagation(t *testing.T) {
// Expected: context is still valid
}

// Execute again - should still work
adapter.SetArgs([]string{"sub"})
err = adapter.ExecuteContext(ctx)
// Execute again - should still work (fresh command tree each time)
err = adapter.ExecuteContext(ctx, []string{"sub"})
require.NoError(t, err)
require.Len(t, receivedContexts, 2, "Second execution should have received context")

Expand All @@ -224,42 +226,43 @@ func Test_WorkflowCmdAdapter_ContextPropagation(t *testing.T) {
// Track which contexts were seen
var receivedContexts []context.Context

// Create a root command with nested subcommands
rootCmd := &cobra.Command{
Use: "root",
}
// Create a command factory that builds a fresh tree on each call
newCommand := func() *cobra.Command {
rootCmd := &cobra.Command{
Use: "root",
}

parentCmd := &cobra.Command{
Use: "parent",
}
parentCmd := &cobra.Command{
Use: "parent",
}

childCmd := &cobra.Command{
Use: "child",
RunE: func(cmd *cobra.Command, args []string) error {
receivedContexts = append(receivedContexts, cmd.Context())
return nil
},
}
childCmd := &cobra.Command{
Use: "child",
RunE: func(cmd *cobra.Command, args []string) error {
receivedContexts = append(receivedContexts, cmd.Context())
return nil
},
}

parentCmd.AddCommand(childCmd)
rootCmd.AddCommand(parentCmd)
parentCmd.AddCommand(childCmd)
rootCmd.AddCommand(parentCmd)
return rootCmd
}

adapter := &workflowCmdAdapter{cmd: rootCmd}
adapter := &workflowCmdAdapter{newCommand: newCommand}

// In production, main.go wraps with context.WithoutCancel.
ctx := context.WithoutCancel(context.Background())
adapter.SetArgs([]string{"parent", "child"})
err := adapter.ExecuteContext(ctx)
err := adapter.ExecuteContext(ctx, []string{"parent", "child"})
require.NoError(t, err)
require.Len(t, receivedContexts, 1)

// Verify context is marked as child action
require.True(t, middleware.IsChildAction(receivedContexts[0]),
"Nested context should be marked as child action")

// Second execution should also work
adapter.SetArgs([]string{"parent", "child"})
err = adapter.ExecuteContext(ctx)
// Second execution should also work (fresh command tree)
err = adapter.ExecuteContext(ctx, []string{"parent", "child"})
require.NoError(t, err)
require.Len(t, receivedContexts, 2)

Expand All @@ -274,4 +277,85 @@ func Test_WorkflowCmdAdapter_ContextPropagation(t *testing.T) {
require.True(t, middleware.IsChildAction(receivedContexts[1]),
"Second nested context should also be marked as child action")
})

t.Run("FreshCommandTreeOnEachExecution", func(t *testing.T) {
// Verify that each ExecuteContext call creates a new command tree,
// ensuring no stale state from previous executions.
var commandTreeInstances []*cobra.Command

newCommand := func() *cobra.Command {
rootCmd := &cobra.Command{
Use: "root",
}
rootCmd.AddCommand(&cobra.Command{
Use: "test",
RunE: func(cmd *cobra.Command, args []string) error {
return nil
},
})
commandTreeInstances = append(commandTreeInstances, rootCmd)
return rootCmd
}

adapter := &workflowCmdAdapter{newCommand: newCommand}
ctx := context.WithoutCancel(context.Background())

err := adapter.ExecuteContext(ctx, []string{"test"})
require.NoError(t, err)

err = adapter.ExecuteContext(ctx, []string{"test"})
require.NoError(t, err)

// Each execution should have created a distinct command tree
require.Len(t, commandTreeInstances, 2, "Factory should have been called twice")
require.NotSame(t, commandTreeInstances[0], commandTreeInstances[1],
"Each execution should use a distinct command tree instance")
})

t.Run("NewRootCmdPreservesMiddlewareChain", func(t *testing.T) {
// Verify that building a real command tree via NewRootCmd preserves
// the full middleware chain (debug, ux, telemetry, error, loginGuard, etc.)
container := ioc.NewNestedContainer(nil)
ctx := context.WithoutCancel(context.Background())
ioc.RegisterInstance(container, ctx)
ioc.RegisterInstance(container, &internal.GlobalCommandOptions{})

rootCmd := NewRootCmd(false, nil, container)

// Verify the command tree is fully built with known subcommands
foundVersion := false
foundProvision := false
foundDeploy := false
for _, child := range rootCmd.Commands() {
switch child.Name() {
case "version":
foundVersion = true
case "provision":
foundProvision = true
case "deploy":
foundDeploy = true
}
}

require.True(t, foundVersion, "version command should be registered")
require.True(t, foundProvision, "provision command should be registered")
require.True(t, foundDeploy, "deploy command should be registered")

// Build a second tree and verify it also has all commands
rootCmd2 := NewRootCmd(false, nil, container)
foundVersion2 := false
foundProvision2 := false
for _, child := range rootCmd2.Commands() {
switch child.Name() {
case "version":
foundVersion2 = true
case "provision":
foundProvision2 = true
}
}

require.True(t, foundVersion2, "second tree: version command should be registered")
require.True(t, foundProvision2, "second tree: provision command should be registered")
require.NotSame(t, rootCmd, rootCmd2, "each NewRootCmd call should produce a distinct instance")
})
}
3 changes: 1 addition & 2 deletions cli/azd/cmd/init.go
Original file line number Diff line number Diff line change
Expand Up @@ -314,8 +314,7 @@ func (i *initAction) Run(ctx context.Context) (*actions.ActionResult, error) {
if deploy {
// Call azd up
startTime := time.Now()
i.azd.SetArgs([]string{"up", "--cwd", azdCtx.ProjectDirectory()})
err := i.azd.ExecuteContext(ctx)
err := i.azd.ExecuteContext(ctx, []string{"up", "--cwd", azdCtx.ProjectDirectory()})
header = "New project initialized! Provision and deploy to Azure was completed in " +
ux.DurationAsText(since(startTime)) + "."
if err != nil {
Expand Down
3 changes: 1 addition & 2 deletions cli/azd/internal/cmd/add/add.go
Original file line number Diff line number Diff line change
Expand Up @@ -359,8 +359,7 @@ func (a *AddAction) Run(ctx context.Context) (*actions.ActionResult, error) {
}

if provisionOption == provision {
a.azd.SetArgs([]string{followUpCmd})
err = a.azd.ExecuteContext(ctx)
err = a.azd.ExecuteContext(ctx, []string{followUpCmd})
if err != nil {
return nil, err
}
Expand Down
27 changes: 9 additions & 18 deletions cli/azd/internal/grpcserver/workflow_service_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,7 @@ func Test_WorkflowService_Run_Success(t *testing.T) {
t.Run("Success", func(t *testing.T) {
testRunner := &TestWorkflowRunner{}
runner := workflow.NewRunner(testRunner, mockContext.Console)
testRunner.On("SetArgs", mock.Anything)
testRunner.On("ExecuteContext", contextType).Return(nil)
testRunner.On("ExecuteContext", contextType, mock.Anything).Return(nil)

service := NewWorkflowService(runner)

Expand All @@ -48,17 +47,15 @@ func Test_WorkflowService_Run_Success(t *testing.T) {
require.NoError(t, err)
require.NotNil(t, resp)

// Verify that the runner's Run method was invoked.
testRunner.AssertCalled(t, "SetArgs", []string{"provision"})
testRunner.AssertCalled(t, "ExecuteContext", contextType)
// Verify that the runner's ExecuteContext was invoked with the correct args.
testRunner.AssertCalled(t, "ExecuteContext", contextType, []string{"provision"})
})

t.Run("Failure", func(t *testing.T) {
expectedErr := errors.New("execution failed")
testRunner := &TestWorkflowRunner{}
runner := workflow.NewRunner(testRunner, mockContext.Console)
testRunner.On("SetArgs", mock.Anything)
testRunner.On("ExecuteContext", contextType).Return(expectedErr)
testRunner.On("ExecuteContext", contextType, mock.Anything).Return(expectedErr)

service := NewWorkflowService(runner)

Expand All @@ -83,9 +80,8 @@ func Test_WorkflowService_Run_Success(t *testing.T) {
require.Error(t, err)
require.Nil(t, resp)

// Verify that the runner's Run method was invoked.
testRunner.AssertCalled(t, "SetArgs", []string{"provision"})
testRunner.AssertCalled(t, "ExecuteContext", contextType)
// Verify that the runner's ExecuteContext was invoked with the correct args.
testRunner.AssertCalled(t, "ExecuteContext", contextType, []string{"provision"})
})
}

Expand All @@ -94,13 +90,8 @@ type TestWorkflowRunner struct {
mock.Mock
}

// Modified SetArgs to use testify/mock.
func (r *TestWorkflowRunner) SetArgs(args []string) {
r.Called(args)
}

// Modified ExecuteContext to use testify/mock.
func (r *TestWorkflowRunner) ExecuteContext(ctx context.Context) error {
ret := r.Called(ctx)
// ExecuteContext implements workflow.AzdCommandRunner using testify/mock.
func (r *TestWorkflowRunner) ExecuteContext(ctx context.Context, args []string) error {
ret := r.Called(ctx, args)
return ret.Error(0)
}
Loading
Loading