Reuse existing azd environment on azd ai agent init reruns

cli/azd/extensions/azure.ai.agents/internal/cmd/init.go

@@ -117,7 +117,7 @@ func runInitFromManifest(
 	}
 
 	// Get or create environment
-	env := getExistingEnvironment(ctx, flags, azdClient)
+	env := getExistingEnvironment(ctx, flags.env, azdClient)
 	if env == nil {
 		fmt.Println("Lets create a new default azd environment for your project.")
 		env, err = createNewEnvironment(ctx, azdClient, flags.env)
@@ -498,15 +498,15 @@ func ensureProject(ctx context.Context, flags *initFlags, azdClient *azdext.AzdC
 	return projectResponse.Project, nil
 }
 
-func getExistingEnvironment(ctx context.Context, flags *initFlags, azdClient *azdext.AzdClient) *azdext.Environment {
+func getExistingEnvironment(ctx context.Context, envName string, azdClient *azdext.AzdClient) *azdext.Environment {
 	var env *azdext.Environment
-	if flags.env == "" {
+	if envName == "" {
 		if envResponse, err := azdClient.Environment().GetCurrent(ctx, &azdext.EmptyRequest{}); err == nil {
 			env = envResponse.Environment
 		}
 	} else {
 		if envResponse, err := azdClient.Environment().Get(ctx, &azdext.GetEnvironmentRequest{
-			Name: flags.env,
+			Name: envName,
 		}); err == nil {
 			env = envResponse.Environment
 		}

cli/azd/extensions/azure.ai.agents/internal/cmd/init_foundry_resources_helpers.go

@@ -18,6 +18,8 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerregistry/armcontainerregistry"
 	"github.com/azure/azure-dev/cli/azd/pkg/azdext"
 	"github.com/azure/azure-dev/cli/azd/pkg/ux"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 // FoundryProjectInfo holds information about a discovered or parsed Foundry project.
@@ -548,6 +550,12 @@ func createNewEnvironment(
 	azdClient *azdext.AzdClient,
 	envName string,
 ) (*azdext.Environment, error) {
+	if envName != "" {
+		if existingEnv := getExistingEnvironment(ctx, envName, azdClient); existingEnv != nil {
+			return existingEnv, nil
+		}
+	}
+
 	envArgs := []string{"env", "new"}
 	if envName != "" {
 		envArgs = append(envArgs, envName)
@@ -567,6 +575,14 @@ func createNewEnvironment(
 		if exterrors.IsCancellation(err) {
 			return nil, exterrors.Cancelled("environment creation was cancelled")
 		}
+		// The workflow may have created the environment on disk before returning
+		// an "already exists" error (e.g. a concurrent process raced, or a previous
+		// partial run left env files behind). Re-fetch so we can reuse it.
+		if envName != "" && status.Code(err) == codes.AlreadyExists {
+			if existingEnv := getExistingEnvironment(ctx, envName, azdClient); existingEnv != nil {
+				return existingEnv, nil
+			}
+		}
 		return nil, exterrors.Dependency(
 			exterrors.CodeEnvironmentCreationFailed,
 			fmt.Sprintf("failed to create new azd environment: %s", err),
@@ -575,7 +591,7 @@ func createNewEnvironment(
 	}
 
 	// Re-fetch the environment after creation
-	env := getExistingEnvironment(ctx, &initFlags{env: envName}, azdClient)
+	env := getExistingEnvironment(ctx, envName, azdClient)
 	if env == nil {
 		return nil, exterrors.Dependency(
 			exterrors.CodeEnvironmentNotFound,

cli/azd/extensions/azure.ai.agents/internal/cmd/init_foundry_resources_helpers_test.go

@@ -4,9 +4,15 @@
 package cmd
 
 import (
+	"context"
+	"net"
 	"testing"
 
+	"github.com/azure/azure-dev/cli/azd/pkg/azdext"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 func TestExtractProjectDetails(t *testing.T) {
@@ -105,3 +111,130 @@ func TestFoundryProjectInfoResourceIdConstruction(t *testing.T) {
 
 	require.Equal(t, originalId, reconstructed)
 }
+
+func TestCreateNewEnvironment_ReturnsExistingNamedEnvironment(t *testing.T) {
+	t.Parallel()
+
+	envServer := &testEnvironmentServiceServer{
+		environments: map[string]*azdext.Environment{
+			"agent-dev": {Name: "agent-dev"},
+		},
+	}
+	workflowServer := &testWorkflowServiceServer{}
+	azdClient := newTestAzdClient(t, envServer, workflowServer)
+
+	env, err := createNewEnvironment(t.Context(), azdClient, "agent-dev")
+
+	require.NoError(t, err)
+	require.Equal(t, "agent-dev", env.Name)
+	require.Equal(t, 0, workflowServer.runCalls)
+}
+
+func TestCreateNewEnvironment_ReusesExistingEnvironmentAfterAlreadyExistsError(t *testing.T) {
+	t.Parallel()
+
+	const envName = "agent-dev"
+
+	envServer := &testEnvironmentServiceServer{
+		environments: make(map[string]*azdext.Environment),
+	}
+	workflowServer := &testWorkflowServiceServer{
+		runErr: status.Error(
+			codes.AlreadyExists,
+			"environment already exists",
+		),
+		runHook: func() {
+			envServer.environments[envName] = &azdext.Environment{Name: envName}
+		},
+	}
+	azdClient := newTestAzdClient(t, envServer, workflowServer)
+
+	env, err := createNewEnvironment(t.Context(), azdClient, envName)
+
+	require.NoError(t, err)
+	require.Equal(t, envName, env.Name)
+	require.Equal(t, 1, workflowServer.runCalls)
+}
+
+type testEnvironmentServiceServer struct {
+	azdext.UnimplementedEnvironmentServiceServer
+	environments map[string]*azdext.Environment
+	current      *azdext.Environment
+}
+
+func (s *testEnvironmentServiceServer) GetCurrent(context.Context, *azdext.EmptyRequest) (*azdext.EnvironmentResponse, error) {
+	if s.current == nil {
+		return nil, status.Error(codes.NotFound, "current environment not found")
+	}
+
+	return &azdext.EnvironmentResponse{Environment: s.current}, nil
+}
+
+func (s *testEnvironmentServiceServer) Get(_ context.Context, req *azdext.GetEnvironmentRequest) (*azdext.EnvironmentResponse, error) {
+	env, ok := s.environments[req.Name]
+	if !ok {
+		return nil, status.Error(codes.NotFound, "environment not found")
+	}
+
+	return &azdext.EnvironmentResponse{Environment: env}, nil
+}
+
+type testWorkflowServiceServer struct {
+	azdext.UnimplementedWorkflowServiceServer
+	runCalls int
+	runErr   error
+	runHook  func()
+}
+
+func (s *testWorkflowServiceServer) Run(context.Context, *azdext.RunWorkflowRequest) (*azdext.EmptyResponse, error) {
+	s.runCalls++
+	if s.runHook != nil {
+		s.runHook()
+	}
+	if s.runErr != nil {
+		return nil, s.runErr
+	}
+
+	return &azdext.EmptyResponse{}, nil
+}
+
+func newTestAzdClient(
+	t *testing.T,
+	envServer azdext.EnvironmentServiceServer,
+	workflowServer azdext.WorkflowServiceServer,
+) *azdext.AzdClient {
+	t.Helper()
+
+	grpcServer := grpc.NewServer()
+	azdext.RegisterEnvironmentServiceServer(grpcServer, envServer)
+	azdext.RegisterWorkflowServiceServer(grpcServer, workflowServer)
+
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+
+	serveErr := make(chan error, 1)
+	go func() {
+		if err := grpcServer.Serve(listener); err != nil {
+			serveErr <- err
+		}
+	}()
+
+	t.Cleanup(func() {
+		grpcServer.Stop()
+		_ = listener.Close()
+		select {
+		case err := <-serveErr:
+			require.ErrorIs(t, err, grpc.ErrServerStopped)
+		default:
+		}
+	})
+
+	azdClient, err := azdext.NewAzdClient(azdext.WithAddress(listener.Addr().String()))
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		azdClient.Close()
+	})
+
+	return azdClient
+}

cli/azd/internal/grpcserver/workflow_service.go

@@ -5,8 +5,10 @@ package grpcserver
 
 import (
 	"context"
+	"errors"
 
 	"github.com/azure/azure-dev/cli/azd/pkg/azdext"
+	"github.com/azure/azure-dev/cli/azd/pkg/environment"
 	"github.com/azure/azure-dev/cli/azd/pkg/workflow"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -39,6 +41,9 @@ func (s *workflowService) Run(ctx context.Context, request *azdext.RunWorkflowRe
 	}
 
 	if err := s.runner.Run(ctx, azdWorkflow); err != nil {
+		if errors.Is(err, environment.ErrExists) {
+			return nil, status.Errorf(codes.AlreadyExists, "failed to run workflow: %v", err)
+		}
 		return nil, status.Errorf(codes.Internal, "failed to run workflow: %v", err)
 	}
 

cli/azd/internal/grpcserver/workflow_service_test.go

@@ -6,13 +6,17 @@ package grpcserver
 import (
 	"context"
 	"errors"
+	"fmt"
 	"testing"
 
 	"github.com/azure/azure-dev/cli/azd/pkg/azdext"
+	"github.com/azure/azure-dev/cli/azd/pkg/environment"
 	"github.com/azure/azure-dev/cli/azd/pkg/workflow"
 	"github.com/azure/azure-dev/cli/azd/test/mocks"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 func Test_WorkflowService_Run_Success(t *testing.T) {
@@ -78,11 +82,36 @@ func Test_WorkflowService_Run_Success(t *testing.T) {
 
 		// Assert
 		require.Error(t, err)
+		require.Equal(t, codes.Internal, status.Code(err))
 		require.Nil(t, resp)
 
 		// Verify that the runner's ExecuteContext was invoked with the correct args.
 		testRunner.AssertCalled(t, "ExecuteContext", contextType, []string{"provision"})
 	})
+
+	t.Run("EnvironmentAlreadyExists", func(t *testing.T) {
+		envExistsErr := fmt.Errorf("creating environment 'myenv': %w", environment.ErrExists)
+		testRunner := &TestWorkflowRunner{}
+		runner := workflow.NewRunner(testRunner, mockContext.Console)
+		testRunner.On("ExecuteContext", contextType, mock.Anything).Return(envExistsErr)
+
+		service := NewWorkflowService(runner)
+
+		req := &azdext.RunWorkflowRequest{
+			Workflow: &azdext.Workflow{
+				Name: "env new",
+				Steps: []*azdext.WorkflowStep{
+					{Command: &azdext.WorkflowCommand{Args: []string{"env", "new", "myenv"}}},
+				},
+			},
+		}
+
+		resp, err := service.Run(*mockContext.Context, req)
+
+		require.Error(t, err)
+		require.Equal(t, codes.AlreadyExists, status.Code(err))
+		require.Nil(t, resp)
+	})
 }
 
 // Updated TestWorkflowRunner using testify/mock.

cli/azd/pkg/environment/manager.go

@@ -170,7 +170,7 @@ func (m *manager) Create(ctx context.Context, spec Spec) (*Environment, error) {
 	case err != nil:
 		return nil, fmt.Errorf("checking for existing environment: %w", err)
 	default:
-		return nil, fmt.Errorf("environment '%s' already exists", spec.Name)
+		return nil, fmt.Errorf("%w: '%s'", ErrExists, spec.Name)
 	}
 
 	env := New(spec.Name)