fix: give each agent space its own LaunchDarkly MCP server

cmd/aws_devops_agent/plan_print_test.go

@@ -84,7 +84,7 @@ func TestTokenPromptAsksWhenInteractive(t *testing.T) {
 	assert.Contains(t, w.calls, "RegisterService")
 	assert.Less(t, strings.Index(stderrOf(cmd), "Paste the token"), strings.Index(stderrOf(cmd), "Setup plan:"))
 	// The plan is built again with the token, so it is current.
-	assert.Contains(t, stderrOf(cmd), "Register MCP server LaunchDarkly at")
+	assert.Contains(t, stderrOf(cmd), "Register a new MCP server for this agent space at")
 	assert.NotContains(t, stderrOf(cmd), "Ask for a token")
 }
 
@@ -178,12 +178,13 @@ func TestSetupNeedsATokenItCannotAsk(t *testing.T) {
 	}
 }
 
-// A registered server needs no token.
+// The server of the agent space needs no token.
 func TestSetupWithoutTokenReusesServer(t *testing.T) {
 	cmd, _ := setupTestCmd(t, "")
 	cmd.SetIn(noReadStdin{t})
-	clients, w := countingClients()
+	clients, w := countingClients(agenttypes.AgentSpace{AgentSpaceId: aws.String("space-1"), Name: aws.String("platform")})
 	agent := clients.Agent.(countingAgent)
+	agent.associations = []agenttypes.Association{{AssociationId: aws.String("assoc-1"), ServiceId: aws.String("mcp-1")}}
 	agent.services = []agenttypes.RegisteredService{{
 		ServiceId:   aws.String("mcp-1"),
 		ServiceType: agenttypes.ServiceMcpServer,
@@ -193,10 +194,12 @@ func TestSetupWithoutTokenReusesServer(t *testing.T) {
 	}}
 	clients.Agent = agent
 
-	require.NoError(t, setupWith(cmd, clients, noTokenOptions(), true, false))
+	opts := noTokenOptions()
+	opts.AgentSpaceName, opts.AgentSpaceID = "", "space-1"
+
+	require.NoError(t, setupWith(cmd, clients, opts, true, false))
 
 	assert.NotContains(t, w.calls, "RegisterService")
-	assert.Contains(t, w.calls, "AssociateService")
 }
 
 func TestSetupPlanNamesAccountAssociation(t *testing.T) {
@@ -429,12 +432,13 @@ func TestSetupFailureNamesTheNewSpace(t *testing.T) {
 	assert.Contains(t, err.Error(), "To use it, run setup again with --agent-space-id space-new")
 }
 
-// A server that appears between the two plans is reused, and setup says
-// that it does not use the pasted token.
+// A server that the space gets between the two plans is reused, and setup
+// says that it does not use the pasted token.
 func TestSetupSaysWhenTheSecondPlanReusesAServer(t *testing.T) {
 	cmd, _ := setupTestCmd(t, "api-token\n")
-	clients, _ := countingClients()
+	clients, _ := countingClients(agenttypes.AgentSpace{AgentSpaceId: aws.String("space-1"), Name: aws.String("platform")})
 	agent := clients.Agent.(countingAgent)
+	agent.associations = []agenttypes.Association{{AssociationId: aws.String("assoc-7"), ServiceId: aws.String("mcp-7")}}
 	agent.services = []agenttypes.RegisteredService{{
 		ServiceId:   aws.String("mcp-7"),
 		ServiceType: agenttypes.ServiceMcpServer,
@@ -444,6 +448,7 @@ func TestSetupSaysWhenTheSecondPlanReusesAServer(t *testing.T) {
 	}}
 	clients.Agent = agent
 	opts := noTokenOptions()
+	opts.AgentSpaceName, opts.AgentSpaceID = "", "space-1"
 	first := awsdevops.SetupPlan{MCP: awsdevops.MCPStep{Action: awsdevops.MCPAsk}}
 
 	plan, err := planWithToken(cmd, clients, &opts, first, true)

cmd/aws_devops_agent/select_test.go

@@ -20,6 +20,8 @@ import (
 type stubAgent struct {
 	spaces   []agenttypes.AgentSpace
 	services []agenttypes.RegisteredService
+	// associations are the associations of every space.
+	associations []agenttypes.Association
 	// unreadTags fails the tag lookup of these space IDs.
 	unreadTags map[string]bool
 	// missing makes GetAgentSpace report these space IDs as not found.
@@ -88,7 +90,7 @@ func (s stubAgent) ListAgentSpaces(context.Context, *devopsagent.ListAgentSpaces
 }
 
 func (s stubAgent) ListAssociations(context.Context, *devopsagent.ListAssociationsInput, ...func(*devopsagent.Options)) (*devopsagent.ListAssociationsOutput, error) {
-	return &devopsagent.ListAssociationsOutput{}, nil
+	return &devopsagent.ListAssociationsOutput{Associations: s.associations}, nil
 }
 
 func (s stubAgent) ListServices(context.Context, *devopsagent.ListServicesInput, ...func(*devopsagent.Options)) (*devopsagent.ListServicesOutput, error) {

cmd/aws_devops_agent/setup.go

@@ -46,11 +46,13 @@ full setup. Pass --agent-space-id to pick a space without the list, or
 --agent-space-name, or type one when setup asks. The run ends with the
 operator app URL, where the agent is used.
 
-The MCP server is connected with --access-token, or with the token in your
-ldcli configuration when you do not pass one. AWS keeps the token and the agent
-acts as its account and role, so a session token from 'ldcli login' stops
-working once it expires. With no token, setup reuses a registered MCP server.
-If none exists, setup asks for a service token in a terminal and checks it
+Each agent space gets its own LaunchDarkly MCP server registration, so each
+space can use a different token. The MCP server is connected with
+--access-token, or with the token in your ldcli configuration when you do not
+pass one. AWS keeps the token and the agent acts as its account and role, so a
+session token from 'ldcli login' stops working once it expires. If the agent
+space already has an MCP server, setup reuses it and does not change its token.
+If it has none, setup asks for a service token in a terminal and checks it
 before it changes AWS. Without a terminal, or with --yes, it stops and asks
 for --access-token.`,
 		Args:   cobra.NoArgs,
@@ -69,9 +71,9 @@ for --access-token.`,
 	cmd.Flags().String(issuerURLFlag, "", "OIDC issuer URL, required when --auth-flow=idp")
 	cmd.Flags().String(idpClientIDFlag, "", "OIDC client ID, required when --auth-flow=idp")
 	cmd.Flags().String(idpClientSecretFlag, "", "OIDC client secret, required when --auth-flow=idp")
-	cmd.Flags().Bool(replaceMCPTokenFlag, false, "Re-register the LaunchDarkly MCP server so it uses the token passed with --access-token")
-	cmd.Flags().String(mcpServiceIDFlag, "", "Registered MCP server to use, instead of the one registered at --mcp-endpoint")
-	cmd.Flags().String(mcpEndpointFlag, "", "MCP endpoint to register and look the server up by. Defaults to "+awsdevops.MCPServerEndpoint)
+	cmd.Flags().Bool(replaceMCPTokenFlag, false, "Give the agent space a new LaunchDarkly MCP server with the token passed with --access-token. Other agent spaces keep their server")
+	cmd.Flags().String(mcpServiceIDFlag, "", "Registered MCP server to connect, instead of the server of the agent space")
+	cmd.Flags().String(mcpEndpointFlag, "", "MCP endpoint to register. Setup reuses only a server at this endpoint. Defaults to "+awsdevops.MCPServerEndpoint)
 
 	cmd.Flags().Bool(yesFlag, false, "Ask no questions and accept every step")
 

internal/awsdevops/awsdevops_test.go

@@ -50,6 +50,9 @@ type fakeAgent struct {
 	registerOutput    *devopsagent.RegisterServiceOutput
 	registerErr       error
 	deregisterErr     error
+	disassociateErr   error
+	// onDisassociate runs before DisassociateService returns.
+	onDisassociate    func()
 	registerInputs    []*devopsagent.RegisterServiceInput
 	associations      []agenttypes.Association
 	agentSpaces       []agenttypes.AgentSpace
@@ -123,6 +126,12 @@ func (f *fakeAgent) DeregisterService(_ context.Context, in *devopsagent.Deregis
 
 func (f *fakeAgent) DisassociateService(_ context.Context, in *devopsagent.DisassociateServiceInput, _ ...func(*devopsagent.Options)) (*devopsagent.DisassociateServiceOutput, error) {
 	f.calls = append(f.calls, "DisassociateService")
+	if f.onDisassociate != nil {
+		f.onDisassociate()
+	}
+	if f.disassociateErr != nil {
+		return nil, f.disassociateErr
+	}
 	f.disassociatedIDs = append(f.disassociatedIDs, aws.ToString(in.AssociationId))
 
 	return &devopsagent.DisassociateServiceOutput{}, nil
@@ -363,8 +372,9 @@ func TestSetupCreatesRolesAgentSpaceAndAssociations(t *testing.T) {
 	assert.Contains(t, iamClient.inlinePolicies[awsdevops.AgentSpaceRoleName], "iam:CreateServiceLinkedRole")
 	assert.Equal(
 		t,
-		// Plan lookup, the pre-write check, and the check before the MCP step.
-		[]string{"ListServices", "ListServices", "CreateAgentSpace", "AssociateService", "EnableOperatorApp", "ListServices", "RegisterService", "AssociateService"},
+		// The plan does not look for the server of a new space. The check
+		// before the MCP step looks, because the space exists by then.
+		[]string{"CreateAgentSpace", "AssociateService", "EnableOperatorApp", "ListServices", "RegisterService", "AssociateService"},
 		agent.calls,
 	)
 }
@@ -501,12 +511,13 @@ func stubLaunchDarkly(t *testing.T, status int) string {
 
 func TestSetupReusesAnExistingMCPServerWithoutAToken(t *testing.T) {
 	agent := &fakeAgent{
-		services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")},
+		services:     []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")},
+		associations: mcpLink("mcp-existing"),
 	}
 
 	result, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{
-		AgentSpaceName: "launchdarkly",
-		AuthFlow:       "iam",
+		AgentSpaceID: "space-1",
+		AuthFlow:     "iam",
 	})
 	require.NoError(t, err)
 
@@ -538,7 +549,7 @@ func TestSetupDisassociatesTheMCPServerBeforeReplacingItsToken(t *testing.T) {
 	}
 
 	result, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{
-		AgentSpaceName:  "launchdarkly",
+		AgentSpaceID:    "space-1",
 		AuthFlow:        "iam",
 		LDBaseURI:       stubLaunchDarkly(t, http.StatusOK),
 		LDAccessToken:   "api-token",
@@ -553,37 +564,6 @@ func TestSetupDisassociatesTheMCPServerBeforeReplacingItsToken(t *testing.T) {
 	assert.Equal(t, "assoc-mcp-1", result.MCPAssociationID)
 }
 
-func TestSetupReconnectsOtherSpacesAfterReplacingTheMCPToken(t *testing.T) {
-	agent := &fakeAgent{
-		agentSpaces: []agenttypes.AgentSpace{
-			{AgentSpaceId: aws.String("space-1")},
-			{AgentSpaceId: aws.String("space-2")},
-		},
-		services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")},
-		associations: []agenttypes.Association{{
-			AssociationId: aws.String("assoc-mcp-existing"),
-			ServiceId:     aws.String("mcp-existing"),
-		}},
-	}
-
-	_, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{
-		AgentSpaceID:    "space-1",
-		AuthFlow:        "iam",
-		LDBaseURI:       stubLaunchDarkly(t, http.StatusOK),
-		LDAccessToken:   "api-token",
-		ReplaceMCPToken: true,
-	})
-	require.NoError(t, err)
-
-	associated := []string{}
-	for _, in := range agent.associationInputs {
-		if aws.ToString(in.ServiceId) == "mcp-1" {
-			associated = append(associated, aws.ToString(in.AgentSpaceId))
-		}
-	}
-	assert.Equal(t, []string{"space-1", "space-2"}, associated)
-}
-
 func TestSetupKeepsTheMCPServerWhenLaunchDarklyRejectsTheNewToken(t *testing.T) {
 	agent := &fakeAgent{
 		services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")},
@@ -607,14 +587,10 @@ func TestSetupKeepsTheMCPServerWhenLaunchDarklyRejectsTheNewToken(t *testing.T)
 	assert.NotContains(t, agent.calls, "RegisterService")
 }
 
-func TestSetupNamesTheSpacesLeftWithoutAnMCPServer(t *testing.T) {
+func TestSetupNamesTheSpaceLeftWithoutAnMCPServer(t *testing.T) {
 	agent := &fakeAgent{
 		registerErr: errors.New("ValidationException: something broke"),
-		agentSpaces: []agenttypes.AgentSpace{
-			{AgentSpaceId: aws.String("space-1")},
-			{AgentSpaceId: aws.String("space-2")},
-		},
-		services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")},
+		services:    []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")},
 		associations: []agenttypes.Association{{
 			AssociationId: aws.String("assoc-mcp-existing"),
 			ServiceId:     aws.String("mcp-existing"),
@@ -629,7 +605,7 @@ func TestSetupNamesTheSpacesLeftWithoutAnMCPServer(t *testing.T) {
 		ReplaceMCPToken: true,
 	})
 
-	assert.ErrorContains(t, err, "agent spaces space-1, space-2 have no LaunchDarkly MCP server")
+	assert.ErrorContains(t, err, "agent spaces space-1 have no LaunchDarkly MCP server")
 	assert.ErrorContains(t, err, "--access-token")
 }
 

internal/awsdevops/lookup.go

@@ -197,29 +197,66 @@ func agentSpaceARN(region, accountID, spaceID string) string {
 	return fmt.Sprintf("arn:aws:aidevops:%s:%s:agentspace/%s", region, accountID, spaceID)
 }
 
-// FindMCPServer returns the ID of the MCP server registered at endpoint, if the
-// account has one. AWS keeps MCP servers at the account level, so a second
-// agent space reuses the existing registration.
-func FindMCPServer(ctx context.Context, clients Clients, endpoint string) (string, error) {
-	id, _, err := findMCPService(ctx, clients, endpoint)
+// MCPServerNameFor is the name of the MCP server that setup registers for an
+// agent space. AWS needs a unique name of letters, digits, "_" and "-", and
+// the ID ties the server to the space.
+func MCPServerNameFor(agentSpaceID string) string {
+	return MCPServerName + "-" + agentSpaceID
+}
 
-	return id, err
+// mcpRef identifies a registered MCP server.
+type mcpRef struct {
+	ID   string
+	Name string
 }
 
-// findMCPService returns the ID and the name of the MCP server registered at
-// endpoint, if the account has one.
-func findMCPService(ctx context.Context, clients Clients, endpoint string) (id, name string, err error) {
+// spaceMCPServers finds the LaunchDarkly MCP servers at endpoint of an agent
+// space. linked is the server that the space is connected to. own is the
+// server with the name that setup gives the space, which can be the same
+// server. AWS keeps MCP servers at the account level, and each one has the
+// token of one LaunchDarkly account, so the server of another space is never
+// a match. A space connected to more than one server is an error, because
+// setup cannot tell which one to use.
+func spaceMCPServers(ctx context.Context, clients Clients, agentSpaceID, endpoint string) (linked, own mcpRef, err error) {
+	// A new space has no ID and no server yet.
+	if agentSpaceID == "" {
+		return mcpRef{}, mcpRef{}, nil
+	}
+	associations, err := listAssociations(ctx, clients, agentSpaceID)
+	if err != nil {
+		return mcpRef{}, mcpRef{}, err
+	}
+	links := make(map[string]bool, len(associations))
+	for _, association := range associations {
+		links[aws.ToString(association.ServiceId)] = true
+	}
 	services, err := listServices(ctx, clients)
 	if err != nil {
-		return "", "", err
+		return mcpRef{}, mcpRef{}, err
 	}
+	var all []string
 	for _, service := range services {
-		if mcpEndpointOf(service) == endpoint {
-			return aws.ToString(service.ServiceId), serviceName(service), nil
+		if mcpEndpointOf(service) != endpoint {
+			continue
+		}
+		ref := mcpRef{ID: aws.ToString(service.ServiceId), Name: serviceName(service)}
+		if links[ref.ID] {
+			linked = ref
+			all = append(all, fmt.Sprintf("%s (%s)", ref.Name, ref.ID))
 		}
+		if own.ID == "" && ref.Name == MCPServerNameFor(agentSpaceID) {
+			own = ref
+		}
+	}
+	if len(all) > 1 {
+		return mcpRef{}, mcpRef{}, fmt.Errorf(
+			"agent space %s is connected to %d LaunchDarkly MCP servers: %s. "+
+				"Disconnect all but one in the AWS console, then run setup again",
+			agentSpaceID, len(all), strings.Join(all, ", "),
+		)
 	}
 
-	return "", "", nil
+	return linked, own, nil
 }
 
 // spacesNamed lists IDs of spaces named name.