feat: turn on elevated actions in aws-devops-agent setup

cmd/aws_devops_agent/confirm_test.go

@@ -71,6 +71,12 @@ func (a countingAgent) GetAgentSpace(_ context.Context, in *devopsagent.GetAgent
 	return &devopsagent.GetAgentSpaceOutput{AgentSpace: &agenttypes.AgentSpace{AgentSpaceId: in.AgentSpaceId}}, nil
 }
 
+func (a countingAgent) UpdateAgentSpace(context.Context, *devopsagent.UpdateAgentSpaceInput, ...func(*devopsagent.Options)) (*devopsagent.UpdateAgentSpaceOutput, error) {
+	a.writes.add("UpdateAgentSpace")
+
+	return &devopsagent.UpdateAgentSpaceOutput{}, nil
+}
+
 func (a countingAgent) RegisterService(context.Context, *devopsagent.RegisterServiceInput, ...func(*devopsagent.Options)) (*devopsagent.RegisterServiceOutput, error) {
 	a.writes.add("RegisterService")
 

cmd/aws_devops_agent/select_test.go

@@ -85,6 +85,10 @@ func (s stubAgent) GetAgentSpace(_ context.Context, in *devopsagent.GetAgentSpac
 	return &devopsagent.GetAgentSpaceOutput{AgentSpace: &agenttypes.AgentSpace{AgentSpaceId: in.AgentSpaceId}}, nil
 }
 
+func (s stubAgent) UpdateAgentSpace(context.Context, *devopsagent.UpdateAgentSpaceInput, ...func(*devopsagent.Options)) (*devopsagent.UpdateAgentSpaceOutput, error) {
+	return nil, nil
+}
+
 func (s stubAgent) ListAgentSpaces(context.Context, *devopsagent.ListAgentSpacesInput, ...func(*devopsagent.Options)) (*devopsagent.ListAgentSpacesOutput, error) {
 	return &devopsagent.ListAgentSpacesOutput{AgentSpaces: s.spaces}, nil
 }

cmd/aws_devops_agent/setup.go

@@ -41,10 +41,11 @@ LaunchDarkly MCP server connection the AWS DevOps Agent needs.
 When the account already has agent spaces, setup asks which one to connect
 LaunchDarkly to. On a space that setup did not make, setup only registers the
 MCP server. On a space that setup made, or with --full-setup, setup runs the
-full setup. Pass --agent-space-id to pick a space without the list, or
---new-agent-space to provision another space. A new space needs a name: pass
---agent-space-name, or type one when setup asks. The run ends with the
-operator app URL, where the agent is used.
+full setup. A full setup turns on elevated actions on the space. Pass
+--agent-space-id to pick a space without the list, or --new-agent-space to
+provision another space. A new space needs a name: pass --agent-space-name,
+or type one when setup asks. The run ends with the operator app URL, where
+the agent is used.
 
 Each agent space gets its own LaunchDarkly MCP server registration, so each
 space can use a different token. The MCP server is connected with

internal/awsdevops/awsdevops.go

@@ -61,6 +61,7 @@ type AgentAPI interface {
 	ListAssociations(context.Context, *devopsagent.ListAssociationsInput, ...func(*devopsagent.Options)) (*devopsagent.ListAssociationsOutput, error)
 	ListServices(context.Context, *devopsagent.ListServicesInput, ...func(*devopsagent.Options)) (*devopsagent.ListServicesOutput, error)
 	RegisterService(context.Context, *devopsagent.RegisterServiceInput, ...func(*devopsagent.Options)) (*devopsagent.RegisterServiceOutput, error)
+	UpdateAgentSpace(context.Context, *devopsagent.UpdateAgentSpaceInput, ...func(*devopsagent.Options)) (*devopsagent.UpdateAgentSpaceOutput, error)
 }
 
 // IAMAPI is the subset of the IAM API this package uses.

internal/awsdevops/awsdevops_test.go

@@ -34,6 +34,10 @@ type fakeAgent struct {
 	tagsErrFor map[string]error
 	// createdTags holds the tags of the created space.
 	createdTags map[string]string
+	// createdPreferences holds the preferences of the created space.
+	createdPreferences map[string]bool
+	// updatedPreferences holds the preferences of each space update.
+	updatedPreferences []map[string]bool
 	// goneSpaces makes GetAgentSpace report not found.
 	goneSpaces      map[string]bool
 	getSpaceErr     error
@@ -108,6 +112,7 @@ func (f *fakeAgent) AssociateService(_ context.Context, in *devopsagent.Associat
 func (f *fakeAgent) CreateAgentSpace(_ context.Context, in *devopsagent.CreateAgentSpaceInput, _ ...func(*devopsagent.Options)) (*devopsagent.CreateAgentSpaceOutput, error) {
 	f.calls = append(f.calls, "CreateAgentSpace")
 	f.createdTags = in.Tags
+	f.createdPreferences = in.Preferences
 
 	return &devopsagent.CreateAgentSpaceOutput{
 		AgentSpace: &agenttypes.AgentSpace{AgentSpaceId: aws.String("space-1")},
@@ -180,18 +185,32 @@ func (f *fakeAgent) GetAgentSpace(_ context.Context, in *devopsagent.GetAgentSpa
 	}
 	name := "launchdarkly"
 	var description *string
+	var preferences map[string]bool
 	for _, space := range f.agentSpaces {
 		if aws.ToString(space.AgentSpaceId) == id && space.Name != nil {
 			name = aws.ToString(space.Name)
 			description = space.Description
+			preferences = space.Preferences
 		}
 	}
 
 	return &devopsagent.GetAgentSpaceOutput{
-		AgentSpace: &agenttypes.AgentSpace{AgentSpaceId: in.AgentSpaceId, Name: aws.String(name), Description: description},
+		AgentSpace: &agenttypes.AgentSpace{
+			AgentSpaceId: in.AgentSpaceId,
+			Name:         aws.String(name),
+			Description:  description,
+			Preferences:  preferences,
+		},
 	}, nil
 }
 
+func (f *fakeAgent) UpdateAgentSpace(_ context.Context, in *devopsagent.UpdateAgentSpaceInput, _ ...func(*devopsagent.Options)) (*devopsagent.UpdateAgentSpaceOutput, error) {
+	f.calls = append(f.calls, "UpdateAgentSpace")
+	f.updatedPreferences = append(f.updatedPreferences, in.Preferences)
+
+	return &devopsagent.UpdateAgentSpaceOutput{}, nil
+}
+
 func (f *fakeAgent) ListAgentSpaces(_ context.Context, in *devopsagent.ListAgentSpacesInput, _ ...func(*devopsagent.Options)) (*devopsagent.ListAgentSpacesOutput, error) {
 	spaces := f.agentSpaces
 	if spaces == nil {
@@ -904,3 +923,69 @@ func TestMCPAssociationClassifiesTheSkillTools(t *testing.T) {
 		assert.Equal(t, agenttypes.ToolClassificationReadOnly, classes[tool], tool)
 	}
 }
+
+func TestSetupTurnsOnElevatedActionsOnANewSpace(t *testing.T) {
+	agent := &fakeAgent{}
+
+	_, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{
+		AgentSpaceName: "launchdarkly",
+		AuthFlow:       "iam",
+		LDAccessToken:  "api-token",
+	})
+	require.NoError(t, err)
+
+	assert.Equal(t, map[string]bool{awsdevops.ElevatedActionsPreference: true}, agent.createdPreferences)
+	assert.NotContains(t, agent.calls, "UpdateAgentSpace")
+}
+
+func TestSetupTurnsOnElevatedActionsOnAReusedSpaceAndKeepsOtherPreferences(t *testing.T) {
+	agent := &fakeAgent{agentSpaces: []agenttypes.AgentSpace{{
+		AgentSpaceId: aws.String("space-existing"),
+		Name:         aws.String("launchdarkly"),
+		Preferences:  map[string]bool{"otherPreference": true},
+	}}}
+	opts := awsdevops.SetupOptions{AgentSpaceName: "launchdarkly", AuthFlow: "iam"}
+
+	plan, err := awsdevops.PlanSetup(context.Background(), newTestClients(agent, newFakeIAM()), opts, true)
+	require.NoError(t, err)
+	assert.Contains(t, plan.Text().AgentSpace, "turn on elevated actions")
+
+	_, err = awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), opts)
+	require.NoError(t, err)
+
+	assert.Equal(t, []map[string]bool{{
+		"otherPreference":                   true,
+		awsdevops.ElevatedActionsPreference: true,
+	}}, agent.updatedPreferences)
+}
+
+func TestSetupLeavesASpaceWithElevatedActionsUnchanged(t *testing.T) {
+	agent := &fakeAgent{agentSpaces: []agenttypes.AgentSpace{{
+		AgentSpaceId: aws.String("space-existing"),
+		Name:         aws.String("launchdarkly"),
+		Preferences:  map[string]bool{awsdevops.ElevatedActionsPreference: true},
+	}}}
+
+	_, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{
+		AgentSpaceName: "launchdarkly",
+		AuthFlow:       "iam",
+	})
+	require.NoError(t, err)
+
+	assert.NotContains(t, agent.calls, "UpdateAgentSpace")
+}
+
+func TestPlanSetupConnectLeavesElevatedActionsUnchanged(t *testing.T) {
+	agent := &fakeAgent{agentSpaces: []agenttypes.AgentSpace{{
+		AgentSpaceId: aws.String("space-existing"),
+		Name:         aws.String("launchdarkly"),
+	}}}
+
+	plan, err := awsdevops.PlanSetup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{
+		AgentSpaceID: "space-existing",
+	}, false)
+	require.NoError(t, err)
+
+	assert.False(t, plan.Space.EnableElevated)
+	assert.NotContains(t, plan.Text().AgentSpace, "elevated")
+}

internal/awsdevops/setup.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"maps"
 	"slices"
 	"strings"
 	"time"
@@ -43,6 +44,9 @@ const (
 	MCPServerEndpoint = "https://mcp.launchdarkly.com/mcp/launchdarkly"
 
 	agentSpaceDescription = "Managed by the LaunchDarkly CLI"
+	// ElevatedActionsPreference lets the agent run actions that change AWS
+	// resources.
+	ElevatedActionsPreference = "elevatedActionsEnabled"
 )
 
 // The agent invokes read-only tools without asking for approval, so anything
@@ -189,17 +193,23 @@ type SpaceStep struct {
 	Action SpaceAction
 	ID     string
 	Name   string
+	// EnableElevated turns on elevated actions on the space.
+	EnableElevated bool
 }
 
 func (s SpaceStep) describe() string {
 	label := fmt.Sprintf("%s (%s)", s.Name, s.ID)
+	elevated := ""
+	if s.EnableElevated {
+		elevated = ", and turn on elevated actions"
+	}
 	switch s.Action {
 	case SpaceCreate:
-		return "Create agent space " + s.Name
+		return "Create agent space " + s.Name + elevated
 	case SpaceConnect:
 		return "Connect to agent space " + label
 	default:
-		return "Reuse agent space " + label
+		return "Reuse agent space " + label + elevated
 	}
 }
 
@@ -587,7 +597,7 @@ func planAgentSpace(ctx context.Context, clients Clients, opts SetupOptions, pro
 		return SpaceStep{}, errors.New("a new agent space needs a name. Pass --agent-space-name")
 	}
 	if id == "" {
-		return SpaceStep{Action: SpaceCreate, Name: opts.AgentSpaceName}, nil
+		return SpaceStep{Action: SpaceCreate, Name: opts.AgentSpaceName, EnableElevated: true}, nil
 	}
 
 	space, err := clients.Agent.GetAgentSpace(ctx, &devopsagent.GetAgentSpaceInput{AgentSpaceId: aws.String(id)})
@@ -597,6 +607,8 @@ func planAgentSpace(ctx context.Context, clients Clients, opts SetupOptions, pro
 	step := SpaceStep{Action: SpaceReuse, ID: id, Name: aws.ToString(space.AgentSpace.Name)}
 	if !provision {
 		step.Action = SpaceConnect
+	} else {
+		step.EnableElevated = !space.AgentSpace.Preferences[ElevatedActionsPreference]
 	}
 
 	return step, nil
@@ -628,6 +640,29 @@ func checkSpace(ctx context.Context, clients Clients, step SpaceStep) error {
 	return nil
 }
 
+// enableElevatedActions turns on elevated actions on a space. An update
+// replaces all the preferences, so it sends the current ones too.
+func enableElevatedActions(ctx context.Context, clients Clients, spaceID string) error {
+	space, err := clients.Agent.GetAgentSpace(ctx, &devopsagent.GetAgentSpaceInput{AgentSpaceId: aws.String(spaceID)})
+	if err != nil {
+		return fmt.Errorf("unable to read agent space %s: %w", spaceID, err)
+	}
+	preferences := maps.Clone(space.AgentSpace.Preferences)
+	if preferences == nil {
+		preferences = map[string]bool{}
+	}
+	preferences[ElevatedActionsPreference] = true
+	_, err = clients.Agent.UpdateAgentSpace(ctx, &devopsagent.UpdateAgentSpaceInput{
+		AgentSpaceId: aws.String(spaceID),
+		Preferences:  preferences,
+	})
+	if err != nil {
+		return fmt.Errorf("unable to turn on elevated actions on agent space %s: %w", spaceID, err)
+	}
+
+	return nil
+}
+
 func planMCPStep(ctx context.Context, clients Clients, opts SetupOptions, spaceID string) (MCPStep, error) {
 	step, err := planMCPAction(ctx, clients, opts, spaceID)
 	step.Endpoint = opts.mcpEndpoint()
@@ -919,11 +954,15 @@ func RunPlan(ctx context.Context, clients Clients, opts SetupOptions, plan Setup
 		if err := checkSpace(ctx, clients, plan.Space); err != nil {
 			return result, err
 		}
-		space, err := clients.Agent.CreateAgentSpace(ctx, &devopsagent.CreateAgentSpaceInput{
+		input := &devopsagent.CreateAgentSpaceInput{
 			Name:        aws.String(plan.Space.Name),
 			Description: aws.String(agentSpaceDescription),
 			Tags:        map[string]string{ManagedTagKey: ManagedTagValue},
-		})
+		}
+		if plan.Space.EnableElevated {
+			input.Preferences = map[string]bool{ElevatedActionsPreference: true}
+		}
+		space, err := clients.Agent.CreateAgentSpace(ctx, input)
 		if err != nil {
 			return result, fmt.Errorf("unable to create the agent space: %w", err)
 		}
@@ -932,6 +971,12 @@ func RunPlan(ctx context.Context, clients Clients, opts SetupOptions, plan Setup
 	} else {
 		result.AgentSpaceID = plan.Space.ID
 		logf("Reused agent space %s (%s)", plan.Space.Name, plan.Space.ID)
+		if plan.Space.EnableElevated {
+			if err := enableElevatedActions(ctx, clients, plan.Space.ID); err != nil {
+				return result, err
+			}
+			logf("Turned on elevated actions on agent space %s", plan.Space.ID)
+		}
 	}
 
 	result.AWSAssociationID, err = FindAWSAssociation(ctx, clients, result.AgentSpaceID, accountID)