Fixes context cancellation issue

cli/azd/cmd/cobra_builder.go

@@ -94,9 +94,7 @@ func (cb *CobraBuilder) configureActionResolver(cmd *cobra.Command, descriptor *
 	}
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
-		// Register root go context that will be used for resolving singleton dependencies
-		ctx := tools.WithInstalledCheckCache(cmd.Context())
-		ioc.RegisterInstance(cb.container, ctx)
+		ctx := cmd.Context()
 
 		// Create new container scope for the current command
 		cmdContainer, err := cb.container.NewScope()

cli/azd/cmd/container.go

@@ -174,7 +174,7 @@ func registerCommonDependencies(container *ioc.NestedContainer) {
 		return writer
 	})
 
-	container.MustRegisterScoped(func(cmd *cobra.Command) internal.EnvFlag {
+	container.MustRegisterScoped(func(ctx context.Context, cmd *cobra.Command) internal.EnvFlag {
 		// The env flag `-e, --environment` is available on most azd commands but not all
 		// This is typically used to override the default environment and is used for bootstrapping other components
 		// such as the azd environment.
@@ -190,7 +190,7 @@ func registerCommonDependencies(container *ioc.NestedContainer) {
 			// If no explicit environment flag was set, but one was provided
 			// in the context, use that instead.
 			// This is used in workflow execution (in `up`) to influence the environment used.
-			if envFlag, ok := cmd.Context().Value(envFlagCtxKey).(internal.EnvFlag); ok {
+			if envFlag, ok := ctx.Value(envFlagCtxKey).(internal.EnvFlag); ok {
 				return envFlag
 			}
 		}
@@ -918,29 +918,9 @@ func (w *workflowCmdAdapter) SetArgs(args []string) {
 // ExecuteContext implements workflow.AzdCommandRunner
 func (w *workflowCmdAdapter) ExecuteContext(ctx context.Context) error {
 	childCtx := middleware.WithChildAction(ctx)
-
-	// CRITICAL FIX: Explicitly set the context on the command before calling ExecuteContext.
-	// When the same command instance is reused across workflow steps, cobra may retain
-	// a stale cancelled context from a previous execution. Setting the context on both
-	// the root command and all subcommands (recursively) ensures the fresh context is used.
-	setContextRecursively(w.cmd, childCtx)
-
 	return w.cmd.ExecuteContext(childCtx)
 }
 
-// setContextRecursively sets the context on a command and all its subcommands recursively
-func setContextRecursively(cmd *cobra.Command, ctx context.Context) {
-	// Always set the context, even if the command hasn't been initialized yet.
-	// cobra.Command.Context() returns context.Background() if no context is set,
-	// so SetContext is safe to call on any command. This ensures all commands
-	// in the tree start with the correct context, preventing stale contexts
-	// from previous executions when commands are reused.
-	cmd.SetContext(ctx)
-	for _, subCmd := range cmd.Commands() {
-		setContextRecursively(subCmd, ctx)
-	}
-}
-
 // ArmClientInitializer is a function definition for all Azure SDK ARM Client
 type ArmClientInitializer[T comparable] func(
 	subscriptionId string,

cli/azd/cmd/container_test.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/azure/azure-dev/cli/azd/cmd/middleware"
 	"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"
@@ -160,11 +161,12 @@ type testConcreteComponent[T comparable] struct {
 }
 
 // Test_WorkflowCmdAdapter_ContextPropagation validates that the workflowCmdAdapter
-// properly sets fresh contexts on both root and subcommands when reused across
-// multiple workflow steps, preventing stale cancelled contexts from affecting
-// subsequent executions.
+// properly marks contexts as child actions when executing subcommands.
+// The main.go entrypoint wraps the root context with context.WithoutCancel,
+// so workflow steps always receive a non-cancellable context.
+// See: https://github.com/Azure/azure-dev/issues/6530
 func Test_WorkflowCmdAdapter_ContextPropagation(t *testing.T) {
-	t.Run("SubcommandReceivesFreshContext", func(t *testing.T) {
+	t.Run("SubcommandReceivesChildActionContext", func(t *testing.T) {
 		// Track which contexts were seen by the subcommand
 		var receivedContexts []context.Context
 
@@ -187,54 +189,38 @@ func Test_WorkflowCmdAdapter_ContextPropagation(t *testing.T) {
 		// Create the adapter
 		adapter := &workflowCmdAdapter{cmd: rootCmd}
 
-		// Simulate first workflow step
-		ctx1, cancel1 := context.WithCancel(context.Background())
+		// 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(ctx1)
+		err := adapter.ExecuteContext(ctx)
 		require.NoError(t, err)
-		require.Len(t, receivedContexts, 1, "First execution should have received context")
+		require.Len(t, receivedContexts, 1, "Execution should have received context")
 
-		// Verify first context is not cancelled
-		select {
-		case <-receivedContexts[0].Done():
-			t.Fatal("First context should not be cancelled during execution")
-		default:
-			// Expected: context is still valid
-		}
-
-		// Cancel the first context (simulating workflow step completion)
-		cancel1()
+		// Verify context is marked as child action
+		require.True(t, middleware.IsChildAction(receivedContexts[0]),
+			"Context should be marked as child action")
 
-		// Verify first context is now cancelled
+		// Verify context is not cancelled (since we used WithoutCancel)
 		select {
 		case <-receivedContexts[0].Done():
-			// Expected: context is cancelled
+			t.Fatal("Context should not be cancelled")
 		default:
-			t.Fatal("First context should be cancelled after cancel1()")
+			// Expected: context is still valid
 		}
 
-		// Simulate second workflow step with a fresh context
-		ctx2 := context.Background()
+		// Execute again - should still work
 		adapter.SetArgs([]string{"sub"})
-		err = adapter.ExecuteContext(ctx2)
+		err = adapter.ExecuteContext(ctx)
 		require.NoError(t, err)
 		require.Len(t, receivedContexts, 2, "Second execution should have received context")
 
-		// CRITICAL TEST: Verify second execution received a valid context,
-		// not the cancelled context from the first execution
-		select {
-		case <-receivedContexts[1].Done():
-			t.Fatal("2nd execution should have received a fresh valid context, not the cancelled one from 1st execution")
-		default:
-			// Expected: second context is valid
-		}
-
-		// Verify the two contexts are different
-		require.NotSame(t, receivedContexts[0], receivedContexts[1],
-			"Second execution should receive a different context than the first")
+		// Both contexts should be marked as child actions
+		require.True(t, middleware.IsChildAction(receivedContexts[1]),
+			"Second context should also be marked as child action")
 	})
 
-	t.Run("NestedSubcommandReceivesFreshContext", func(t *testing.T) {
+	t.Run("NestedSubcommandReceivesChildActionContext", func(t *testing.T) {
 		// Track which contexts were seen
 		var receivedContexts []context.Context
 
@@ -260,29 +246,32 @@ func Test_WorkflowCmdAdapter_ContextPropagation(t *testing.T) {
 
 		adapter := &workflowCmdAdapter{cmd: rootCmd}
 
-		// First execution
-		ctx1, cancel1 := context.WithCancel(context.Background())
+		// In production, main.go wraps with context.WithoutCancel.
+		ctx := context.WithoutCancel(context.Background())
 		adapter.SetArgs([]string{"parent", "child"})
-		err := adapter.ExecuteContext(ctx1)
+		err := adapter.ExecuteContext(ctx)
 		require.NoError(t, err)
 		require.Len(t, receivedContexts, 1)
 
-		// Cancel first context
-		cancel1()
+		// Verify context is marked as child action
+		require.True(t, middleware.IsChildAction(receivedContexts[0]),
+			"Nested context should be marked as child action")
 
-		// Second execution with fresh context
-		ctx2 := context.Background()
+		// Second execution should also work
 		adapter.SetArgs([]string{"parent", "child"})
-		err = adapter.ExecuteContext(ctx2)
+		err = adapter.ExecuteContext(ctx)
 		require.NoError(t, err)
 		require.Len(t, receivedContexts, 2)
 
-		// Verify second execution got a valid context
+		// Verify second execution got a valid context marked as child
 		select {
 		case <-receivedContexts[1].Done():
-			t.Fatal("Nested subcommand should have received a fresh valid context in second execution")
+			t.Fatal("Nested subcommand should have received a valid context")
 		default:
 			// Expected: context is valid
 		}
+
+		require.True(t, middleware.IsChildAction(receivedContexts[1]),
+			"Second nested context should also be marked as child action")
 	})
 }

cli/azd/cmd/middleware/debug.go

@@ -38,7 +38,7 @@ func NewDebugMiddleware(options *Options, console input.Console) Middleware {
 // a debugger before continuing invocation of the action
 func (m *DebugMiddleware) Run(ctx context.Context, next NextFn) (*actions.ActionResult, error) {
 	// Don't run for sub actions
-	if m.options.IsChildAction(ctx) {
+	if IsChildAction(ctx) {
 		return next(ctx)
 	}
 

cli/azd/cmd/middleware/error.go

@@ -77,7 +77,7 @@ func (e *ErrorMiddleware) Run(ctx context.Context, next NextFn) (*actions.Action
 	// Stop the spinner always to un-hide cursor
 	e.console.StopSpinner(ctx, "", input.Step)
 
-	if err == nil || e.options.IsChildAction(ctx) {
+	if err == nil || IsChildAction(ctx) {
 		return actionResult, err
 	}
 

cli/azd/cmd/middleware/extensions.go

@@ -57,7 +57,7 @@ func NewExtensionsMiddleware(
 
 func (m *ExtensionsMiddleware) Run(ctx context.Context, next NextFn) (*actions.ActionResult, error) {
 	// Extensions were already started in the root parent command
-	if m.options.IsChildAction(ctx) {
+	if IsChildAction(ctx) {
 		return next(ctx)
 	}
 

cli/azd/cmd/middleware/hooks.go

@@ -55,7 +55,7 @@ func NewHooksMiddleware(
 // Runs the Hooks middleware
 func (m *HooksMiddleware) Run(ctx context.Context, next NextFn) (*actions.ActionResult, error) {
 	// Validate hooks and display any warnings
-	if !m.options.IsChildAction(ctx) {
+	if !IsChildAction(ctx) {
 		if err := m.validateHooks(ctx, m.projectConfig); err != nil {
 			return nil, fmt.Errorf("failed validating hooks, %w", err)
 		}

cli/azd/cmd/middleware/middleware.go

@@ -39,11 +39,6 @@ type Options struct {
 	Annotations map[string]string
 }
 
-func (o *Options) IsChildAction(ctx context.Context) bool {
-	value, ok := ctx.Value(childActionKey).(bool)
-	return ok && value
-}
-
 // Sets the container to be used for resolving middleware components
 func (o *Options) WithContainer(container *ioc.NestedContainer) {
 	o.container = container
@@ -143,3 +138,10 @@ func (r *MiddlewareRunner) Use(name string, resolveFn any) error {
 func WithChildAction(ctx context.Context) context.Context {
 	return context.WithValue(ctx, childActionKey, true)
 }
+
+// IsChildAction checks if the given context was created by WithChildAction.
+// This is used to determine if a command is being executed as part of a workflow step.
+func IsChildAction(ctx context.Context) bool {
+	value, ok := ctx.Value(childActionKey).(bool)
+	return ok && value
+}

cli/azd/cmd/middleware/telemetry.go

@@ -65,7 +65,7 @@ func (m *TelemetryMiddleware) Run(ctx context.Context, next NextFn) (*actions.Ac
 
 	log.Printf("TraceID: %s", span.SpanContext().TraceID())
 
-	if !m.options.IsChildAction(ctx) {
+	if !IsChildAction(ctx) {
 		// Set the command name as a baggage item on the span context.
 		// This allow inner actions to have command name attached.
 		spanCtx = tracing.SetBaggageInContext(

cli/azd/cmd/middleware/ux.go

@@ -34,7 +34,7 @@ func NewUxMiddleware(options *Options, console input.Console, featuresManager *a
 
 func (m *UxMiddleware) Run(ctx context.Context, next NextFn) (*actions.ActionResult, error) {
 	// Don't run for sub actions
-	if m.options.IsChildAction(ctx) {
+	if IsChildAction(ctx) {
 		return next(ctx)
 	}
 

cli/azd/cmd/root.go

@@ -133,9 +133,18 @@ func NewRootCmd(
 			// we can just remove the entire project folder afterwards.
 			// In practical execution, this wouldn't affect much, since the CLI is exiting.
 			if prevDir != "" {
-				return os.Chdir(prevDir)
+				if err := os.Chdir(prevDir); err != nil {
+					return err
+				}
 			}
 
+			// Clear context from the executed command to prevent stale context issues.
+			// This is important because singletons (like credential providers) may cache
+			// contexts from earlier commands, leading to "context canceled" errors when
+			// those contexts are cancelled but the singleton is reused.
+			// See: https://github.com/Azure/azure-dev/issues/6530
+			clearCommandContext(cmd)
+
 			return nil
 		},
 		SilenceUsage:      true,
@@ -543,3 +552,14 @@ func getCmdRootHelpCommands(cmd *cobra.Command) (result string) {
 	}
 	return strings.Join(paragraph, "\n")
 }
+
+// clearCommandContext recursively clears the context from the command and its children.
+// This prevents stale context issues where a cancelled context might be reused by singletons
+// or cached components. See: https://github.com/Azure/azure-dev/issues/6530
+func clearCommandContext(cmd *cobra.Command) {
+	//nolint:staticcheck // SA1012: Intentionally passing nil to clear the context and reset to cobra's default state
+	cmd.SetContext(nil)
+	for _, child := range cmd.Commands() {
+		clearCommandContext(child)
+	}
+}

cli/azd/internal/grpcserver/project_service_test.go

@@ -48,13 +48,12 @@ func Test_ProjectService_NoProject(t *testing.T) {
 	})
 
 	// Create mock GitHub CLI.
-	ghCli, err := github.NewGitHubCli(*mockContext.Context, mockContext.Console, mockContext.CommandRunner)
-	require.NoError(t, err)
+	ghCli := github.NewGitHubCli(mockContext.Console, mockContext.CommandRunner)
 
 	// Create the service with ImportManager.
 	importManager := project.NewImportManager(&project.DotNetImporter{})
 	service := NewProjectService(lazyAzdContext, lazyEnvManager, lazyProjectConfig, importManager, ghCli)
-	_, err = service.Get(*mockContext.Context, &azdext.EmptyRequest{})
+	_, err := service.Get(*mockContext.Context, &azdext.EmptyRequest{})
 	require.Error(t, err)
 }
 
@@ -105,8 +104,7 @@ func Test_ProjectService_Flow(t *testing.T) {
 	require.NoError(t, err)
 
 	// Create mock GitHub CLI.
-	ghCli, err := github.NewGitHubCli(*mockContext.Context, mockContext.Console, mockContext.CommandRunner)
-	require.NoError(t, err)
+	ghCli := github.NewGitHubCli(mockContext.Console, mockContext.CommandRunner)
 
 	// Create the service with ImportManager.
 	importManager := project.NewImportManager(&project.DotNetImporter{})
@@ -154,8 +152,7 @@ func Test_ProjectService_AddService(t *testing.T) {
 	lazyProjectConfig := lazy.From(&projectConfig)
 
 	// Create mock GitHub CLI.
-	ghCli, err := github.NewGitHubCli(*mockContext.Context, mockContext.Console, mockContext.CommandRunner)
-	require.NoError(t, err)
+	ghCli := github.NewGitHubCli(mockContext.Console, mockContext.CommandRunner)
 
 	// Create the project service with ImportManager.
 	importManager := project.NewImportManager(&project.DotNetImporter{})

cli/azd/internal/scaffold/scaffold_test.go

@@ -256,7 +256,8 @@ func TestExecInfra(t *testing.T) {
 			}
 
 			ctx := context.Background()
-			cli, err := bicep.NewCli(ctx, mockinput.NewMockConsole(), exec.NewCommandRunner(nil))
+			cli := bicep.NewCli(mockinput.NewMockConsole(), exec.NewCommandRunner(nil))
+			err = cli.EnsureInstalled(ctx)
 			require.NoError(t, err)
 
 			res, err := cli.Build(ctx, filepath.Join(dir, "main.bicep"))

cli/azd/main.go

@@ -33,6 +33,7 @@ import (
 	"github.com/azure/azure-dev/cli/azd/pkg/oneauth"
 	"github.com/azure/azure-dev/cli/azd/pkg/osutil"
 	"github.com/azure/azure-dev/cli/azd/pkg/output"
+	"github.com/azure/azure-dev/cli/azd/pkg/tools"
 	"github.com/blang/semver/v4"
 	"github.com/mattn/go-colorable"
 	"github.com/spf13/pflag"
@@ -67,6 +68,11 @@ func main() {
 	go fetchLatestVersion(latest)
 
 	rootContainer := ioc.NewNestedContainer(nil)
+
+	ctx = context.WithoutCancel(ctx)
+	ctx = tools.WithInstalledCheckCache(ctx)
+
+	// Register the context for singleton resolution
 	ioc.RegisterInstance(rootContainer, ctx)
 
 	// Execute command with auto-installation support for extensions