From bf4064f8a1ec78aceeb686adfc2bf7c047433c4f Mon Sep 17 00:00:00 2001 From: Christina Chu Date: Wed, 28 Oct 2026 11:52:43 -0700 Subject: [PATCH 1/6] feat: register one OAuth MCP server for each AWS account aws-devops-agent setup registered one LaunchDarkly MCP server for each agent space, with a LaunchDarkly access token. AWS keeps the credentials on the registration, so a different token needed a different server. Now setup registers one server for the AWS account, with OAuth. Setup opens the approval URL that AWS sends, waits until AWS registers the server, and associates the server with the agent space. Another agent space reuses the server and needs no browser step. This change removes the token path: the token prompt, the token check, --replace-mcp-token, and the server name for each space. Co-Authored-By: Claude Opus 5.5 --- cmd/aws_devops_agent/confirm_test.go | 14 +- cmd/aws_devops_agent/flags_test.go | 19 - cmd/aws_devops_agent/main_test.go | 11 - cmd/aws_devops_agent/plan_print_test.go | 148 +---- cmd/aws_devops_agent/prompt.go | 16 - cmd/aws_devops_agent/prompt_test.go | 9 - cmd/aws_devops_agent/select_test.go | 8 - cmd/aws_devops_agent/setup.go | 124 +--- internal/awsdevops/awsdevops.go | 2 - internal/awsdevops/awsdevops_test.go | 206 +------ internal/awsdevops/export_test.go | 14 +- internal/awsdevops/lookup.go | 72 +-- internal/awsdevops/lookup_test.go | 98 +-- internal/awsdevops/main_test.go | 21 - internal/awsdevops/plan_test.go | 38 +- ...{regions_token_test.go => regions_test.go} | 20 - internal/awsdevops/runplan_test.go | 570 +++--------------- internal/awsdevops/setup.go | 510 ++++------------ internal/awsdevops/setup_test.go | 45 +- internal/awsdevops/token.go | 63 -- 20 files changed, 311 insertions(+), 1697 deletions(-) delete mode 100644 internal/awsdevops/main_test.go rename internal/awsdevops/{regions_token_test.go => regions_test.go} (75%) delete mode 100644 internal/awsdevops/token.go diff --git a/cmd/aws_devops_agent/confirm_test.go b/cmd/aws_devops_agent/confirm_test.go index 41e543770..950584a84 100644 --- a/cmd/aws_devops_agent/confirm_test.go +++ b/cmd/aws_devops_agent/confirm_test.go @@ -43,18 +43,6 @@ func (a countingAgent) CreateAgentSpace(context.Context, *devopsagent.CreateAgen return &devopsagent.CreateAgentSpaceOutput{AgentSpace: &agenttypes.AgentSpace{AgentSpaceId: aws.String("space-new")}}, nil } -func (a countingAgent) DeregisterService(context.Context, *devopsagent.DeregisterServiceInput, ...func(*devopsagent.Options)) (*devopsagent.DeregisterServiceOutput, error) { - a.writes.add("DeregisterService") - - return &devopsagent.DeregisterServiceOutput{}, nil -} - -func (a countingAgent) DisassociateService(context.Context, *devopsagent.DisassociateServiceInput, ...func(*devopsagent.Options)) (*devopsagent.DisassociateServiceOutput, error) { - a.writes.add("DisassociateService") - - return &devopsagent.DisassociateServiceOutput{}, nil -} - func (a countingAgent) EnableOperatorApp(context.Context, *devopsagent.EnableOperatorAppInput, ...func(*devopsagent.Options)) (*devopsagent.EnableOperatorAppOutput, error) { a.writes.add("EnableOperatorApp") @@ -161,7 +149,7 @@ func setupTestCmd(t *testing.T, stdin string) (*cobra.Command, *bytes.Buffer) { } func newSpaceOptions() awsdevops.SetupOptions { - return awsdevops.SetupOptions{AgentSpaceName: "ld-devops-agent", AuthFlow: "iam", LDAccessToken: "api-token"} + return awsdevops.SetupOptions{AgentSpaceName: "ld-devops-agent", AuthFlow: "iam"} } // A json run in a terminal asks too, because the question goes to stderr. diff --git a/cmd/aws_devops_agent/flags_test.go b/cmd/aws_devops_agent/flags_test.go index e3656583b..f96ff79c9 100644 --- a/cmd/aws_devops_agent/flags_test.go +++ b/cmd/aws_devops_agent/flags_test.go @@ -6,11 +6,9 @@ import ( "github.com/aws/aws-sdk-go-v2/aws" agenttypes "github.com/aws/aws-sdk-go-v2/service/devopsagent/types" - "github.com/spf13/viper" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "github.com/launchdarkly/ldcli/cmd/cliflags" "github.com/launchdarkly/ldcli/internal/analytics" ) @@ -55,7 +53,6 @@ func spaceIDRun(t *testing.T, flags map[string]string) []string { }) opts, err := setupOptions(cmd) require.NoError(t, err) - opts.LDAccessToken = "api-token" require.NoError(t, setupWith(cmd, clients, opts, true, false)) @@ -154,19 +151,3 @@ func TestSetupOptionsRecordsAuthFlow(t *testing.T) { require.NoError(t, err) assert.True(t, opts.FullSetup) } - -// Setup does not ask for the replacement token, so the flag needs one. -func TestReplaceMCPTokenNeedsAToken(t *testing.T) { - viper.Set(cliflags.AccessTokenFlag, "") - t.Cleanup(func() { viper.Set(cliflags.AccessTokenFlag, "") }) - cmd := NewSetupCmd(nil) - require.NoError(t, cmd.Flags().Set(replaceMCPTokenFlag, "true")) - - _, err := setupOptions(cmd) - assert.EqualError(t, err, "--replace-mcp-token needs a token. Pass --access-token") - - viper.Set(cliflags.AccessTokenFlag, "api-token") - opts, err := setupOptions(cmd) - require.NoError(t, err) - assert.True(t, opts.ReplaceMCPToken) -} diff --git a/cmd/aws_devops_agent/main_test.go b/cmd/aws_devops_agent/main_test.go index 7a158abd6..df69fa75e 100644 --- a/cmd/aws_devops_agent/main_test.go +++ b/cmd/aws_devops_agent/main_test.go @@ -1,7 +1,6 @@ package awsdevopsagent import ( - "context" "fmt" "os" "strconv" @@ -9,19 +8,9 @@ import ( "testing" "github.com/spf13/cobra" - - "github.com/launchdarkly/ldcli/internal/awsdevops" ) -// Never call the real LaunchDarkly API. func TestMain(m *testing.M) { - awsdevops.CheckToken = func(ctx context.Context, baseURI, token string) error { - if baseURI == "" { - return nil - } - - return awsdevops.ValidateAccessToken(ctx, baseURI, token) - } newChooser = lineChooser os.Exit(m.Run()) } diff --git a/cmd/aws_devops_agent/plan_print_test.go b/cmd/aws_devops_agent/plan_print_test.go index 814f5fe01..84aed818b 100644 --- a/cmd/aws_devops_agent/plan_print_test.go +++ b/cmd/aws_devops_agent/plan_print_test.go @@ -67,68 +67,6 @@ func TestSetupPlanNamesRolesAndOperatorApp(t *testing.T) { assert.Contains(t, out.String(), "operator app") } -func noTokenOptions() awsdevops.SetupOptions { - opts := newSpaceOptions() - opts.LDAccessToken = "" - - return opts -} - -// The token prompt comes before the plan question. -func TestTokenPromptAsksWhenInteractive(t *testing.T) { - cmd, _ := setupTestCmd(t, "api-token\ny\n") - clients, w := countingClients() - - require.NoError(t, setupWith(cmd, clients, noTokenOptions(), true, true)) - - 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 a new MCP server for this agent space at") - assert.NotContains(t, stderrOf(cmd), "Ask for a token") -} - -// A rejected token stops setup before any AWS write. -func TestSetupChecksAPastedTokenFirst(t *testing.T) { - saved := awsdevops.CheckToken - t.Cleanup(func() { awsdevops.CheckToken = saved }) - awsdevops.CheckToken = func(context.Context, string, string) error { return awsdevops.ErrInvalidAccessToken } - cmd, _ := setupTestCmd(t, "expired-token\ny\n") - clients, w := countingClients() - - err := setupWith(cmd, clients, noTokenOptions(), true, true) - - require.ErrorIs(t, err, awsdevops.ErrInvalidAccessToken) - assert.Empty(t, w.calls) - assert.NotContains(t, stderrOf(cmd), "Setup plan:") -} - -// Enter skips the MCP server, and the end names the next step. -func TestSetupSkippedTokenNamesTheNextStep(t *testing.T) { - cmd, _ := setupTestCmd(t, "\ny\n") - clients, w := countingClients() - - require.NoError(t, setupWith(cmd, clients, noTokenOptions(), true, true)) - - assert.NotContains(t, w.calls, "RegisterService") - assert.Contains(t, stderrOf(cmd), "Skip the MCP server of this agent space") - assert.NotContains(t, stderrOf(cmd), "Ask for a token") - assert.Contains(t, stdoutOf(cmd), "The LaunchDarkly MCP server is not connected") -} - -// Without a token, a connect run has nothing to do. -func TestSetupConnectWithoutATokenStops(t *testing.T) { - cmd, _ := setupTestCmd(t, "2\n\n") - clients, w := twoSimilarSpaces() - - err := setupWith(cmd, clients, noTokenOptions(), true, true) - - require.Error(t, err) - assert.Contains(t, err.Error(), "nothing to do on agent space LaunchDarkly (space-2)") - assert.Empty(t, w.calls) - assert.NotContains(t, stderrOf(cmd), "Continue?") -} - // Progress goes to stderr, so a json run shows it, and stdout stays JSON. func TestSetupJSONShowsProgressOnStderr(t *testing.T) { cmd, _ := setupTestCmd(t, "") @@ -149,35 +87,6 @@ func TestPlanNamesTheAWSDefaultProfile(t *testing.T) { assert.Equal(t, "corp", profileLabel(cmd)) } -// No token, no prompt: stop first. -func TestSetupNeedsATokenItCannotAsk(t *testing.T) { - cases := map[string]struct { - plaintext, interactive, yes bool - }{ - "no terminal": {plaintext: true}, - "json": {}, - "yes in terminal": {plaintext: true, interactive: true, yes: true}, - } - for name, c := range cases { - t.Run(name, func(t *testing.T) { - cmd, out := setupTestCmd(t, "") - cmd.SetIn(noReadStdin{t}) - if c.yes { - require.NoError(t, cmd.Flags().Set(yesFlag, "true")) - } - clients, w := countingClients() - - err := setupWith(cmd, clients, noTokenOptions(), c.plaintext, c.interactive) - - require.Error(t, err) - assert.Contains(t, err.Error(), "--access-token") - assert.Contains(t, err.Error(), awsdevops.AccessTokenURL("")) - assert.Empty(t, w.calls) - assert.NotContains(t, out.String(), "Setup plan") - }) - } -} - // The server of the agent space needs no token. func TestSetupWithoutTokenReusesServer(t *testing.T) { cmd, _ := setupTestCmd(t, "") @@ -194,7 +103,7 @@ func TestSetupWithoutTokenReusesServer(t *testing.T) { }} clients.Agent = agent - opts := noTokenOptions() + opts := newSpaceOptions() opts.AgentSpaceName, opts.AgentSpaceID = "", "space-1" require.NoError(t, setupWith(cmd, clients, opts, true, false)) @@ -370,7 +279,6 @@ func TestSetupChosenSpaceWithFullSetup(t *testing.T) { clients, w := twoSimilarSpaces() opts, err := setupOptions(cmd) require.NoError(t, err) - opts.LDAccessToken = "api-token" require.NoError(t, setupWith(cmd, clients, opts, true, true)) @@ -389,7 +297,6 @@ func TestSetupChosenSpaceWithOnlyAuthFlowStops(t *testing.T) { clients, w := twoSimilarSpaces() opts, err := setupOptions(cmd) require.NoError(t, err) - opts.LDAccessToken = "api-token" err = setupWith(cmd, clients, opts, true, true) @@ -450,56 +357,3 @@ func TestSetupFailureNamesTheNewSpace(t *testing.T) { // The same flags stop on the name of the new space, so give the ID. assert.Contains(t, err.Error(), "To use it, run setup again with --agent-space-id space-new") } - -// 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(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, - AdditionalServiceDetails: &agenttypes.AdditionalServiceDetailsMemberMcpserver{ - Value: agenttypes.RegisteredMCPServerDetails{Name: aws.String("Prod LD"), Endpoint: aws.String(awsdevops.MCPServerEndpoint)}, - }, - }} - 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) - - require.NoError(t, err) - assert.Equal(t, awsdevops.MCPReuse, plan.MCP.Action) - assert.Contains(t, stderrOf(cmd), "MCP server Prod LD (mcp-7) was registered after the first plan") - assert.Contains(t, stderrOf(cmd), "--access-token and --replace-mcp-token") -} - -// Without a name from AWS, the warning gives only the ID, which identifies -// the server, and does not guess a name. -func TestSetupNamesANamelessServerByItsID(t *testing.T) { - cmd, _ := setupTestCmd(t, "api-token\n") - 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, - AdditionalServiceDetails: &agenttypes.AdditionalServiceDetailsMemberMcpserver{ - Value: agenttypes.RegisteredMCPServerDetails{Endpoint: aws.String(awsdevops.MCPServerEndpoint)}, - }, - }} - 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) - - require.NoError(t, err) - require.Equal(t, awsdevops.MCPReuse, plan.MCP.Action) - assert.Contains(t, stderrOf(cmd), "MCP server mcp-7 was registered after the first plan") -} diff --git a/cmd/aws_devops_agent/prompt.go b/cmd/aws_devops_agent/prompt.go index c1dfc1214..564050093 100644 --- a/cmd/aws_devops_agent/prompt.go +++ b/cmd/aws_devops_agent/prompt.go @@ -68,22 +68,6 @@ func readYes(in io.Reader) bool { return false } -// readSecret reads one line without echoing it, so a pasted credential stays -// off the screen. -func readSecret(in io.Reader) string { - file, ok := in.(*os.File) - if !ok || !term.IsTerminal(int(file.Fd())) { - return readLine(in) - } - - typed, err := term.ReadPassword(int(file.Fd())) - if err != nil { - return "" - } - - return strings.TrimSpace(string(typed)) -} - // canPrompt reports a terminal on stdin and stderr. The prompts go to stderr, // so a run with stderr sent to a file does not wait for an answer. func canPrompt() bool { diff --git a/cmd/aws_devops_agent/prompt_test.go b/cmd/aws_devops_agent/prompt_test.go index f5a71047b..d028c4383 100644 --- a/cmd/aws_devops_agent/prompt_test.go +++ b/cmd/aws_devops_agent/prompt_test.go @@ -2,21 +2,12 @@ package awsdevopsagent import ( "os" - "strings" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -func TestReadSecretTrimsThePastedToken(t *testing.T) { - assert.Equal(t, "api-token", readSecret(strings.NewReader("api-token\n"))) -} - -func TestReadSecretReadsAnEmptyLine(t *testing.T) { - assert.Empty(t, readSecret(strings.NewReader("\n"))) -} - // A file is not wrapped, so each prompt must leave the next lines unread. func TestReadLineLeavesLaterLinesInAFile(t *testing.T) { r, w, err := os.Pipe() diff --git a/cmd/aws_devops_agent/select_test.go b/cmd/aws_devops_agent/select_test.go index 3d0555c90..75bdee8d7 100644 --- a/cmd/aws_devops_agent/select_test.go +++ b/cmd/aws_devops_agent/select_test.go @@ -36,14 +36,6 @@ func (s stubAgent) CreateAgentSpace(context.Context, *devopsagent.CreateAgentSpa return nil, nil } -func (s stubAgent) DeregisterService(context.Context, *devopsagent.DeregisterServiceInput, ...func(*devopsagent.Options)) (*devopsagent.DeregisterServiceOutput, error) { - return nil, nil -} - -func (s stubAgent) DisassociateService(context.Context, *devopsagent.DisassociateServiceInput, ...func(*devopsagent.Options)) (*devopsagent.DisassociateServiceOutput, error) { - return nil, nil -} - func (s stubAgent) GetOperatorApp(context.Context, *devopsagent.GetOperatorAppInput, ...func(*devopsagent.Options)) (*devopsagent.GetOperatorAppOutput, error) { return nil, &agenttypes.ResourceNotFoundException{} } diff --git a/cmd/aws_devops_agent/setup.go b/cmd/aws_devops_agent/setup.go index 3ff70ec21..ef707d558 100644 --- a/cmd/aws_devops_agent/setup.go +++ b/cmd/aws_devops_agent/setup.go @@ -3,12 +3,11 @@ package awsdevopsagent import ( "errors" "fmt" - "os" "strings" agenttypes "github.com/aws/aws-sdk-go-v2/service/devopsagent/types" + "github.com/pkg/browser" "github.com/spf13/cobra" - "github.com/spf13/viper" "github.com/launchdarkly/ldcli/cmd/cliflags" resourcescmd "github.com/launchdarkly/ldcli/cmd/resources" @@ -25,7 +24,6 @@ const ( issuerURLFlag = "issuer-url" idpClientIDFlag = "idp-client-id" idpClientSecretFlag = "idp-client-secret" - replaceMCPTokenFlag = "replace-mcp-token" mcpServiceIDFlag = "mcp-service-id" mcpEndpointFlag = "mcp-endpoint" yesFlag = "yes" @@ -47,15 +45,12 @@ 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 ---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.`, +The AWS account gets one LaunchDarkly MCP server, which all agent spaces use. +If the account has no server, setup registers one with OAuth. It opens a +browser, where you approve access to LaunchDarkly, and waits until AWS +registers the server. AWS can ask you to sign in to the AWS console first. +The agent then acts as the LaunchDarkly member who approved access, in every +agent space. If the agent space already has an MCP server, setup keeps it.`, Args: cobra.NoArgs, PreRun: trackRun(analyticsTrackerFn), RunE: runSetup, @@ -73,8 +68,7 @@ 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, "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(mcpServiceIDFlag, "", "Registered MCP server to connect, instead of the server of the AWS account") 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") @@ -99,13 +93,6 @@ func runSetup(cmd *cobra.Command, args []string) error { if warning := awsdevops.CheckAWSCLIVersion(cmd.Context()); warning != "" { _, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Warning: %s\n", warning) } - if opts.LDAccessToken != "" && accessTokenFromConfig(cmd) { - _, _ = fmt.Fprintln( - cmd.ErrOrStderr(), - "Using the access token from your ldcli configuration for the MCP server. "+ - "Pass --access-token with a service token if that one expires", - ) - } // The prompts go to stderr, so a json run in a terminal can ask too. return setupWith(cmd, clients, opts, plaintext, canPrompt()) } @@ -143,15 +130,6 @@ func setupWith( authFlowFlag, plan.Space.Name, plan.Space.ID, fullSetupFlag, authFlowFlag, ) } - if plan.MCP.Action == awsdevops.MCPAsk && !ask { - return fmt.Errorf("create a service token at %s, then pass --access-token", awsdevops.AccessTokenURL(opts.LDBaseURI)) - } - if plan.MCP.Action == awsdevops.MCPAsk { - plan, err = planWithToken(cmd, clients, &opts, plan, provision) - if err != nil { - return err - } - } // The plan goes to stderr, so a json run shows it too. if err := confirmSetup(cmd, plan, ask); err != nil { return err @@ -177,13 +155,6 @@ func setupWith( return printJSON(cmd, result) } - if result.MCPServiceID == "" { - _, _ = fmt.Fprintf(cmd.OutOrStdout(), - "\nThe LaunchDarkly MCP server is not connected. Create a service token at %s, "+ - "then run setup again with --access-token\n", - awsdevops.AccessTokenURL(opts.LDBaseURI)) - } - // Only a connect run leaves the app off, on a space that setup did not // make. The full setup would change that space, so point to the console. if !result.OperatorAppEnabled { @@ -619,65 +590,6 @@ func agentSpaceLabels(spaces []awsdevops.AgentSpaceSummary) []string { return labels } -// planWithToken asks for the token before any AWS write. With a token, it plans -// again, so the plan that the user confirms is current. Without one, the plan -// skips the MCP server, and a connect run has nothing to do. -func planWithToken( - cmd *cobra.Command, - clients awsdevops.Clients, - opts *awsdevops.SetupOptions, - plan awsdevops.SetupPlan, - provision bool, -) (awsdevops.SetupPlan, error) { - token, err := askToken(cmd, opts.LDBaseURI) - switch { - case err != nil: - return plan, err - case token != "": - opts.LDAccessToken = token - again, err := awsdevops.PlanSetup(cmd.Context(), clients, *opts, provision) - if err == nil && again.MCP.Action == awsdevops.MCPReuse { - // The first plan found no server, so one appeared since then. - // Without a name from AWS, only the ID identifies it. - server := again.MCP.ServiceID - if again.MCP.Name != "" { - server = fmt.Sprintf("%s (%s)", again.MCP.Name, again.MCP.ServiceID) - } - _, _ = fmt.Fprintf(cmd.ErrOrStderr(), - "MCP server %s was registered after the first plan, so setup reuses it and does not "+ - "use the token that you pasted. To replace its token, run setup with --%s and --%s.\n", - server, cliflags.AccessTokenFlag, replaceMCPTokenFlag) - } - - return again, err - case !provision: - return plan, fmt.Errorf( - "setup did not change AWS: without a token, there is nothing to do on agent space %s (%s). Pass --access-token", - plan.Space.Name, plan.Space.ID) - } - plan.MCP.Action = awsdevops.MCPSkip - - return plan, nil -} - -// askToken asks for a LaunchDarkly service token and checks it. An empty -// answer skips the MCP server. -func askToken(cmd *cobra.Command, baseURI string) (string, error) { - out := cmd.ErrOrStderr() - _, _ = fmt.Fprintf(out, "\nCreate a LaunchDarkly service token at:\n %s\n\n", awsdevops.AccessTokenURL(baseURI)) - _, _ = fmt.Fprint(out, "Paste the token, or press Enter to skip: ") - token := readSecret(cmd.InOrStdin()) - _, _ = fmt.Fprintln(out) - if token == "" { - return "", nil - } - if err := awsdevops.CheckToken(cmd.Context(), baseURI, token); err != nil { - return "", fmt.Errorf("setup did not change AWS: %w", err) - } - - return token, nil -} - func setupOptions(cmd *cobra.Command) (awsdevops.SetupOptions, error) { opts := awsdevops.SetupOptions{ AgentSpaceID: mustString(cmd, agentSpaceIDFlag), @@ -690,11 +602,9 @@ func setupOptions(cmd *cobra.Command) (awsdevops.SetupOptions, error) { IdpClientID: mustString(cmd, idpClientIDFlag), IdpClientSecret: mustString(cmd, idpClientSecretFlag), NewAgentSpace: mustBool(cmd, newAgentSpaceFlag), - LDBaseURI: viper.GetString(cliflags.BaseURIFlag), - LDAccessToken: viper.GetString(cliflags.AccessTokenFlag), - ReplaceMCPToken: mustBool(cmd, replaceMCPTokenFlag), MCPServiceID: mustString(cmd, mcpServiceIDFlag), MCPEndpoint: mustString(cmd, mcpEndpointFlag), + OpenURL: browser.OpenURL, } idcSet := opts.IdcInstanceARN != "" @@ -719,25 +629,9 @@ func setupOptions(cmd *cobra.Command) (awsdevops.SetupOptions, error) { default: return opts, fmt.Errorf("--%s must be iam, idc or idp", authFlowFlag) } - // The replacement token comes from the flag, LD_ACCESS_TOKEN, or the ldcli - // configuration. Setup does not ask for one, so without one the flag does - // nothing. - if opts.ReplaceMCPToken && opts.LDAccessToken == "" { - return opts, fmt.Errorf("--%s needs a token. Pass --%s", replaceMCPTokenFlag, cliflags.AccessTokenFlag) - } - return opts, nil } -// accessTokenFromConfig reports whether the token came from the ldcli -// configuration rather than this run, where `ldcli login` may have written a -// session token that AWS would keep past its expiry. -func accessTokenFromConfig(cmd *cobra.Command) bool { - return !cmd.Flags().Changed(cliflags.AccessTokenFlag) && - os.Getenv("LD_ACCESS_TOKEN") == "" && - viper.GetString(cliflags.AccessTokenFlag) != "" -} - // mustString and friends read flags this package declares itself, so a lookup // error is a programming error rather than user input. func mustString(cmd *cobra.Command, name string) string { diff --git a/internal/awsdevops/awsdevops.go b/internal/awsdevops/awsdevops.go index ae234bd8b..23d0f2621 100644 --- a/internal/awsdevops/awsdevops.go +++ b/internal/awsdevops/awsdevops.go @@ -51,8 +51,6 @@ var ( type AgentAPI interface { AssociateService(context.Context, *devopsagent.AssociateServiceInput, ...func(*devopsagent.Options)) (*devopsagent.AssociateServiceOutput, error) CreateAgentSpace(context.Context, *devopsagent.CreateAgentSpaceInput, ...func(*devopsagent.Options)) (*devopsagent.CreateAgentSpaceOutput, error) - DeregisterService(context.Context, *devopsagent.DeregisterServiceInput, ...func(*devopsagent.Options)) (*devopsagent.DeregisterServiceOutput, error) - DisassociateService(context.Context, *devopsagent.DisassociateServiceInput, ...func(*devopsagent.Options)) (*devopsagent.DisassociateServiceOutput, error) EnableOperatorApp(context.Context, *devopsagent.EnableOperatorAppInput, ...func(*devopsagent.Options)) (*devopsagent.EnableOperatorAppOutput, error) GetOperatorApp(context.Context, *devopsagent.GetOperatorAppInput, ...func(*devopsagent.Options)) (*devopsagent.GetOperatorAppOutput, error) ListTagsForResource(context.Context, *devopsagent.ListTagsForResourceInput, ...func(*devopsagent.Options)) (*devopsagent.ListTagsForResourceOutput, error) diff --git a/internal/awsdevops/awsdevops_test.go b/internal/awsdevops/awsdevops_test.go index 8882e8d11..d5a8d7001 100644 --- a/internal/awsdevops/awsdevops_test.go +++ b/internal/awsdevops/awsdevops_test.go @@ -4,8 +4,6 @@ import ( "context" "errors" "fmt" - "net/http" - "net/http/httptest" "os" "path/filepath" "strconv" @@ -53,15 +51,9 @@ type fakeAgent struct { associationInputs []*devopsagent.AssociateServiceInput 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 - disassociatedIDs []string - deregisteredIDs []string associateFailures int // failAssociateFor fails the MCP server association of these spaces. failAssociateFor map[string]bool @@ -119,29 +111,6 @@ func (f *fakeAgent) CreateAgentSpace(_ context.Context, in *devopsagent.CreateAg }, nil } -func (f *fakeAgent) DeregisterService(_ context.Context, in *devopsagent.DeregisterServiceInput, _ ...func(*devopsagent.Options)) (*devopsagent.DeregisterServiceOutput, error) { - f.calls = append(f.calls, "DeregisterService") - if f.deregisterErr != nil { - return nil, f.deregisterErr - } - f.deregisteredIDs = append(f.deregisteredIDs, aws.ToString(in.ServiceId)) - - return &devopsagent.DeregisterServiceOutput{}, nil -} - -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 -} - func (f *fakeAgent) GetOperatorApp(_ context.Context, in *devopsagent.GetOperatorAppInput, _ ...func(*devopsagent.Options)) (*devopsagent.GetOperatorAppOutput, error) { f.calls = append(f.calls, "GetOperatorApp") if f.operatorAppErr != nil { @@ -375,7 +344,6 @@ func TestSetupCreatesRolesAgentSpaceAndAssociations(t *testing.T) { result, err := awsdevops.Setup(context.Background(), newTestClients(agent, iamClient), awsdevops.SetupOptions{ AgentSpaceName: "launchdarkly", AuthFlow: "iam", - LDAccessToken: "api-token", }) require.NoError(t, err) @@ -391,9 +359,13 @@ func TestSetupCreatesRolesAgentSpaceAndAssociations(t *testing.T) { assert.Contains(t, iamClient.inlinePolicies[awsdevops.AgentSpaceRoleName], "iam:CreateServiceLinkedRole") assert.Equal( t, - // 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"}, + // The plan and the first check look only for the server of the + // account, because a new space has no server. The check before the + // MCP step also looks for the server of the space, which exists by then. + []string{ + "ListServices", "ListServices", "CreateAgentSpace", "AssociateService", "EnableOperatorApp", + "ListServices", "ListServices", "RegisterService", "AssociateService", + }, agent.calls, ) } @@ -426,7 +398,7 @@ func TestSetupRetriesTheAccountAssociationUntilTheRoleIsAssumable(t *testing.T) require.NoError(t, err) assert.Equal(t, "assoc-aws", result.AWSAssociationID) - assert.Len(t, agent.associationInputs, 2) + assert.Len(t, awsAssociationInputs(agent), 2) } func TestSetupReusesAgentSpaceAndAssociations(t *testing.T) { @@ -447,7 +419,19 @@ func TestSetupReusesAgentSpaceAndAssociations(t *testing.T) { assert.Equal(t, "space-existing", result.AgentSpaceID) assert.Equal(t, "assoc-aws", result.AWSAssociationID) assert.NotContains(t, agent.calls, "CreateAgentSpace") - assert.Empty(t, agent.associationInputs) + assert.Empty(t, awsAssociationInputs(agent)) +} + +// awsAssociationInputs lists the AWS account associations that setup sent. +func awsAssociationInputs(agent *fakeAgent) []*devopsagent.AssociateServiceInput { + var inputs []*devopsagent.AssociateServiceInput + for _, in := range agent.associationInputs { + if aws.ToString(in.ServiceId) == "aws" { + inputs = append(inputs, in) + } + } + + return inputs } func TestSetupTrustsAgentSpacesInEveryRegion(t *testing.T) { @@ -513,22 +497,7 @@ func mcpService(serviceID, name, endpoint string) agenttypes.RegisteredService { } } -// stubLaunchDarkly stands in for the LaunchDarkly API the access token is -// checked against, and returns its base URI. -func stubLaunchDarkly(t *testing.T, status int) string { - t.Helper() - - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - assert.Equal(t, "/api/v2/caller-identity", r.URL.Path) - assert.Equal(t, "api-token", r.Header.Get("Authorization")) - w.WriteHeader(status) - })) - t.Cleanup(server.Close) - - return server.URL -} - -func TestSetupReusesAnExistingMCPServerWithoutAToken(t *testing.T) { +func TestSetupReusesAnExistingMCPServer(t *testing.T) { agent := &fakeAgent{ services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")}, associations: mcpLink("mcp-existing"), @@ -545,113 +514,6 @@ func TestSetupReusesAnExistingMCPServerWithoutAToken(t *testing.T) { assert.NotContains(t, agent.calls, "RegisterService") } -func TestSetupWithoutATokenLeavesTheMCPServerUnconnected(t *testing.T) { - agent := &fakeAgent{} - - result, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{ - AgentSpaceName: "launchdarkly", - AuthFlow: "iam", - }) - require.NoError(t, err) - - assert.Empty(t, result.MCPServiceID) - assert.NotContains(t, agent.calls, "RegisterService") -} - -func TestSetupDisassociatesTheMCPServerBeforeReplacingItsToken(t *testing.T) { - agent := &fakeAgent{ - services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")}, - associations: []agenttypes.Association{{ - AssociationId: aws.String("assoc-mcp-existing"), - ServiceId: aws.String("mcp-existing"), - }}, - } - - result, 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) - - assert.Equal(t, []string{"assoc-mcp-existing"}, agent.disassociatedIDs) - assert.Equal(t, []string{"mcp-existing"}, agent.deregisteredIDs) - assert.Less(t, indexOf(agent.calls, "DisassociateService"), indexOf(agent.calls, "DeregisterService")) - assert.Equal(t, "mcp-1", result.MCPServiceID) - assert.Equal(t, "assoc-mcp-1", result.MCPAssociationID) -} - -func TestSetupKeepsTheMCPServerWhenLaunchDarklyRejectsTheNewToken(t *testing.T) { - agent := &fakeAgent{ - 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.StatusUnauthorized), - LDAccessToken: "api-token", - ReplaceMCPToken: true, - }) - - assert.ErrorIs(t, err, awsdevops.ErrInvalidAccessToken) - assert.Empty(t, agent.disassociatedIDs) - assert.Empty(t, agent.deregisteredIDs) - assert.NotContains(t, agent.calls, "RegisterService") -} - -func TestSetupNamesTheSpaceLeftWithoutAnMCPServer(t *testing.T) { - agent := &fakeAgent{ - registerErr: errors.New("ValidationException: something broke"), - 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, - }) - - assert.ErrorContains(t, err, "agent spaces space-1 have no LaunchDarkly MCP server") - assert.ErrorContains(t, err, "--access-token") -} - -func TestSetupIgnoresAnMCPServerAtAnotherEndpoint(t *testing.T) { - agent := &fakeAgent{ - services: []agenttypes.RegisteredService{ - mcpService("mcp-self-hosted", awsdevops.MCPServerName, "https://mcp.example.internal/mcp"), - }, - associations: []agenttypes.Association{{ - AssociationId: aws.String("assoc-mcp-self-hosted"), - ServiceId: aws.String("mcp-self-hosted"), - }}, - } - - result, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{ - AgentSpaceID: "space-1", - AuthFlow: "iam", - LDAccessToken: "api-token", - ReplaceMCPToken: true, - }) - require.NoError(t, err) - - assert.Equal(t, "mcp-1", result.MCPServiceID) - assert.Empty(t, agent.deregisteredIDs) - assert.Empty(t, agent.disassociatedIDs) -} - func TestSetupUsesTheMCPServerItIsPointedAt(t *testing.T) { agent := &fakeAgent{ services: []agenttypes.RegisteredService{ @@ -675,10 +537,9 @@ func TestSetupRegistersAtTheEndpointItIsGiven(t *testing.T) { agent := &fakeAgent{} _, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{ - AgentSpaceID: "space-1", - AuthFlow: "iam", - LDAccessToken: "api-token", - MCPEndpoint: "https://mcp.example.internal/mcp", + AgentSpaceID: "space-1", + AuthFlow: "iam", + MCPEndpoint: "https://mcp.example.internal/mcp", }) require.NoError(t, err) @@ -713,16 +574,6 @@ func TestStatusFollowsEveryListPage(t *testing.T) { assert.Len(t, status.AgentSpaces[0].Associations, 2) } -func indexOf(calls []string, name string) int { - for i, call := range calls { - if call == name { - return i - } - } - - return -1 -} - func TestStatusReportsMissingRoles(t *testing.T) { agent := &fakeAgent{ associations: []agenttypes.Association{{AssociationId: aws.String("assoc-1"), ServiceId: aws.String("aws")}}, @@ -805,11 +656,6 @@ func TestManualStepURLsPointAtRegistrationPages(t *testing.T) { "https://eu-west-1.console.aws.amazon.com/aidevops/home?region=eu-west-1", awsdevops.ConsoleURL("eu-west-1"), ) - assert.Equal( - t, - "https://app.launchdarkly.com/settings/authorization", - awsdevops.AccessTokenURL(""), - ) } func TestNewClientsRejectsUnsupportedRegion(t *testing.T) { @@ -894,7 +740,6 @@ func TestMCPAssociationClassifiesTheSkillTools(t *testing.T) { _, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{ AgentSpaceName: "launchdarkly", AuthFlow: "iam", - LDAccessToken: "api-token", }) require.NoError(t, err) @@ -930,7 +775,6 @@ func TestSetupTurnsOnElevatedActionsOnANewSpace(t *testing.T) { _, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{ AgentSpaceName: "launchdarkly", AuthFlow: "iam", - LDAccessToken: "api-token", }) require.NoError(t, err) diff --git a/internal/awsdevops/export_test.go b/internal/awsdevops/export_test.go index 1bf7e3171..ce46da366 100644 --- a/internal/awsdevops/export_test.go +++ b/internal/awsdevops/export_test.go @@ -1,6 +1,9 @@ package awsdevops -import "context" +import ( + "context" + "time" +) // Setup plans a full setup and runs it, for tests. func Setup(ctx context.Context, clients Clients, opts SetupOptions) (SetupResult, error) { @@ -26,3 +29,12 @@ func ConnectExisting(ctx context.Context, clients Clients, opts SetupOptions) (S func MCPTools() []string { return append(append([]string{}, mcpReadOnlyTools...), mcpMutativeTools...) } + +// SetOAuthPollInterval changes how often setup looks for the OAuth server, +// for tests. It returns a function that restores the old interval. +func SetOAuthPollInterval(d time.Duration) func() { + old := oauthPollInterval + oauthPollInterval = d + + return func() { oauthPollInterval = old } +} diff --git a/internal/awsdevops/lookup.go b/internal/awsdevops/lookup.go index a8efbdaed..102de076d 100644 --- a/internal/awsdevops/lookup.go +++ b/internal/awsdevops/lookup.go @@ -197,19 +197,28 @@ func agentSpaceARN(region, accountID, spaceID string) string { return fmt.Sprintf("arn:aws:aidevops:%s:%s:agentspace/%s", region, accountID, spaceID) } -// mcpServerIDLength is the length to truncate the agent space ID to in the -// server name. -const mcpServerIDLength = 12 - -// 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 end of the ID ties the server to the space. -func MCPServerNameFor(agentSpaceID string) string { - if len(agentSpaceID) > mcpServerIDLength { - agentSpaceID = agentSpaceID[len(agentSpaceID)-mcpServerIDLength:] +// accountMCPServer finds the MCP server of the AWS account: a server at +// endpoint that signs in with OAuth. A server with a token is never a match, +// even with the name MCPServerName, because its token belongs to one +// LaunchDarkly account and not to the AWS account. +func accountMCPServer(ctx context.Context, clients Clients, endpoint string) (mcpRef, error) { + services, err := listServices(ctx, clients) + if err != nil { + return mcpRef{}, err + } + for _, service := range services { + if mcpEndpointOf(service) == endpoint && usesOAuth(service) { + return mcpRef{ID: aws.ToString(service.ServiceId), Name: serviceName(service)}, nil + } } - return MCPServerName + "-" + agentSpaceID + return mcpRef{}, nil +} + +func usesOAuth(service agenttypes.RegisteredService) bool { + details, ok := service.AdditionalServiceDetails.(*agenttypes.AdditionalServiceDetailsMemberMcpserver) + + return ok && details.Value.AuthorizationMethod == agenttypes.MCPServerAuthorizationMethodOauth3lo } // mcpRef identifies a registered MCP server. @@ -218,21 +227,18 @@ type mcpRef struct { Name string } -// 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) { +// linkedMCPServer finds the LaunchDarkly MCP server at endpoint that an agent +// space is connected to. Setup keeps that server, so a space that someone +// connected by hand keeps its server. A space connected to more than one +// server is an error, because setup cannot tell which one to use. +func linkedMCPServer(ctx context.Context, clients Clients, agentSpaceID, endpoint string) (mcpRef, error) { // A new space has no ID and no server yet. if agentSpaceID == "" { - return mcpRef{}, mcpRef{}, nil + return mcpRef{}, nil } associations, err := listAssociations(ctx, clients, agentSpaceID) if err != nil { - return mcpRef{}, mcpRef{}, err + return mcpRef{}, err } links := make(map[string]bool, len(associations)) for _, association := range associations { @@ -240,31 +246,29 @@ func spaceMCPServers(ctx context.Context, clients Clients, agentSpaceID, endpoin } services, err := listServices(ctx, clients) if err != nil { - return mcpRef{}, mcpRef{}, err + return mcpRef{}, err } - var all []string + var ( + linked mcpRef + all []string + ) for _, service := range services { - 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 mcpEndpointOf(service) != endpoint || !links[ref.ID] { + continue } + linked = ref + all = append(all, fmt.Sprintf("%s (%s)", ref.Name, ref.ID)) } if len(all) > 1 { - return mcpRef{}, mcpRef{}, fmt.Errorf( + return 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 linked, own, nil + return linked, nil } // spacesNamed lists IDs of spaces named name. diff --git a/internal/awsdevops/lookup_test.go b/internal/awsdevops/lookup_test.go index bded76370..431f119a7 100644 --- a/internal/awsdevops/lookup_test.go +++ b/internal/awsdevops/lookup_test.go @@ -3,7 +3,6 @@ package awsdevops_test import ( "context" "errors" - "net/http" "testing" "github.com/aws/aws-sdk-go-v2/aws" @@ -21,52 +20,6 @@ func TestSetupReusesTheMCPServerOfTheAgentSpace(t *testing.T) { spaceLinks: map[string][]agenttypes.Association{"space-1": mcpLink("mcp-1")}, } - result, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{ - AgentSpaceID: "space-1", - AuthFlow: "iam", - LDAccessToken: "api-token", - }) - require.NoError(t, err) - - assert.NotContains(t, agent.calls, "RegisterService") - assert.Equal(t, "mcp-1", result.MCPServiceID) - assert.Equal(t, "assoc-mcp-1", result.MCPAssociationID) -} - -// The server of another space has the token of that space, so a space never -// uses it. -func TestSetupRegistersAServerBesideTheServerOfAnotherSpace(t *testing.T) { - agent := &fakeAgent{ - agentSpaces: []agenttypes.AgentSpace{{AgentSpaceId: aws.String("space-1")}, {AgentSpaceId: aws.String("space-2")}}, - services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-other")}, - spaceLinks: map[string][]agenttypes.Association{"space-2": mcpLink("mcp-other")}, - } - - result, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{ - AgentSpaceID: "space-1", - AuthFlow: "iam", - LDAccessToken: "api-token", - }) - require.NoError(t, err) - - assert.Equal(t, "mcp-1", result.MCPServiceID) - require.Len(t, agent.registerInputs, 1) - details, ok := agent.registerInputs[0].ServiceDetails.(*agenttypes.ServiceDetailsMemberMcpserver) - require.True(t, ok) - assert.Equal(t, awsdevops.MCPServerNameFor("space-1"), aws.ToString(details.Value.Name)) - assert.Empty(t, agent.disassociatedIDs) -} - -// A run that registered the server and failed to connect it leaves a server -// with the name of the space. The next run connects that server. -func TestSetupReusesTheServerItRegisteredForTheSpace(t *testing.T) { - agent := &fakeAgent{ - services: []agenttypes.RegisteredService{ - mcpService("mcp-other", awsdevops.MCPServerNameFor("space-2"), awsdevops.MCPServerEndpoint), - mcpService("mcp-own", awsdevops.MCPServerNameFor("space-1"), awsdevops.MCPServerEndpoint), - }, - } - result, err := awsdevops.Setup(context.Background(), newTestClients(agent, newFakeIAM()), awsdevops.SetupOptions{ AgentSpaceID: "space-1", AuthFlow: "iam", @@ -74,52 +27,8 @@ func TestSetupReusesTheServerItRegisteredForTheSpace(t *testing.T) { require.NoError(t, err) assert.NotContains(t, agent.calls, "RegisterService") - assert.Equal(t, "mcp-own", result.MCPServiceID) -} - -func TestSetupReplacesTheMCPServerTokenWhenAsked(t *testing.T) { - agent := &fakeAgent{ - services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-old")}, - spaceLinks: map[string][]agenttypes.Association{"space-1": mcpLink("mcp-old")}, - } - - result, 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) - - assert.Equal(t, []string{"mcp-old"}, agent.deregisteredIDs) - assert.Equal(t, "mcp-1", result.MCPServiceID) -} - -// A replacement on a shared server changes only this space. The other spaces -// keep the server and its token. -func TestSetupReplaceLeavesASharedServerToTheOtherSpaces(t *testing.T) { - agent := replaceAcrossSpaces() - - result, 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) - - assert.Empty(t, agent.deregisteredIDs) - assert.Equal(t, []string{"assoc-mcp-old"}, agent.disassociatedIDs) assert.Equal(t, "mcp-1", result.MCPServiceID) - var 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"}, associated) + assert.Equal(t, "assoc-mcp-1", result.MCPAssociationID) } const space1ARN = "arn:aws:aidevops:us-east-1:" + testAccountID + ":agentspace/space-1" @@ -157,9 +66,8 @@ func TestConnectExistingOnlyConnectsTheMCPServer(t *testing.T) { iamClient := newFakeIAM() result, err := awsdevops.ConnectExisting(context.Background(), newTestClients(agent, iamClient), awsdevops.SetupOptions{ - AgentSpaceID: "space-2", - AuthFlow: "iam", - LDAccessToken: "api-token", + AgentSpaceID: "space-2", + AuthFlow: "iam", }) require.NoError(t, err) diff --git a/internal/awsdevops/main_test.go b/internal/awsdevops/main_test.go deleted file mode 100644 index 5d75a1840..000000000 --- a/internal/awsdevops/main_test.go +++ /dev/null @@ -1,21 +0,0 @@ -package awsdevops_test - -import ( - "context" - "os" - "testing" - - "github.com/launchdarkly/ldcli/internal/awsdevops" -) - -// Never call the real LaunchDarkly API. -func TestMain(m *testing.M) { - awsdevops.CheckToken = func(ctx context.Context, baseURI, token string) error { - if baseURI == "" { - return nil - } - - return awsdevops.ValidateAccessToken(ctx, baseURI, token) - } - os.Exit(m.Run()) -} diff --git a/internal/awsdevops/plan_test.go b/internal/awsdevops/plan_test.go index 104055151..1c7caca7b 100644 --- a/internal/awsdevops/plan_test.go +++ b/internal/awsdevops/plan_test.go @@ -274,17 +274,15 @@ func TestPlanTextOfAConnectPlan(t *testing.T) { assert.Equal(t, "No change", text.Association) assert.Equal(t, []string{"No change"}, text.Roles) assert.Equal(t, tc.want, text.OperatorApp) - assert.Equal(t, "Register a new MCP server for this agent space at "+awsdevops.MCPServerEndpoint, text.MCPServer) + assert.Equal(t, + "Register MCP server LaunchDarkly-OAuth for the AWS account at "+awsdevops.MCPServerEndpoint+ + ". You approve access to LaunchDarkly in a browser", + text.MCPServer, + ) }) } } -func TestPlanTextOfASkippedMCPServer(t *testing.T) { - plan := awsdevops.SetupPlan{MCP: awsdevops.MCPStep{Action: awsdevops.MCPSkip}} - - assert.Equal(t, "Skip the MCP server of this agent space, because no token was given", plan.Text().MCPServer) -} - // The plan names the server by its name in AWS, which the customer can change. func TestPlanNamesTheReusedMCPServerByItsName(t *testing.T) { for name, id := range map[string]string{"connected to the space": "", "given by ID": "mcp-1"} { @@ -327,20 +325,6 @@ func TestPlanConnectWithoutANameAsksForTheID(t *testing.T) { assert.EqualError(t, err, "no agent space to connect to. Pass --agent-space-id") } -// A replace plan on a shared server names the spaces that keep it. -func TestPlanReplaceNamesEachSpace(t *testing.T) { - opts := runPlanOptions() - opts.ReplaceMCPToken = true - - plan, err := awsdevops.PlanSetup(context.Background(), newTestClients(replaceAcrossSpaces(), newFakeIAM()), opts, true) - - require.NoError(t, err) - assert.Equal(t, - "Register a new MCP server for this agent space at "+awsdevops.MCPServerEndpoint+ - ". MCP server LaunchDarkly (mcp-old) keeps its token for space-3", - plan.Text().MCPServer) -} - // Two spaces with one name leave no safe choice. func TestPlanStopsOnTwoSpacesWithTheName(t *testing.T) { agent := &fakeAgent{agentSpaces: []agenttypes.AgentSpace{ @@ -430,7 +414,7 @@ func TestFindAWSAssociationMatchesOnlyTheMonitorOfTheAccount(t *testing.T) { } func runPlanOptions() awsdevops.SetupOptions { - return awsdevops.SetupOptions{AgentSpaceName: "launchdarkly", AuthFlow: "iam", LDAccessToken: "api-token"} + return awsdevops.SetupOptions{AgentSpaceName: "launchdarkly", AuthFlow: "iam"} } func reuseWithApp(app *devopsagent.GetOperatorAppOutput) *fakeAgent { @@ -449,7 +433,7 @@ func idcApp() *devopsagent.GetOperatorAppOutput { } } -func replaceAcrossSpaces() *fakeAgent { +func sharedAcrossSpaces() *fakeAgent { return &fakeAgent{ agentSpaces: []agenttypes.AgentSpace{ namedSpace("space-1", "launchdarkly"), namedSpace("space-2", "other"), namedSpace("space-3", "third"), @@ -459,14 +443,6 @@ func replaceAcrossSpaces() *fakeAgent { } } -// replaceOnOneSpace is replaceAcrossSpaces with mcp-old on space-1 only. -func replaceOnOneSpace() *fakeAgent { - agent := replaceAcrossSpaces() - agent.spaceLinks = map[string][]agenttypes.Association{"space-1": mcpLink("mcp-old")} - - return agent -} - func mcpLink(serviceID string) []agenttypes.Association { return []agenttypes.Association{{AssociationId: aws.String("assoc-" + serviceID), ServiceId: aws.String(serviceID)}} } diff --git a/internal/awsdevops/regions_token_test.go b/internal/awsdevops/regions_test.go similarity index 75% rename from internal/awsdevops/regions_token_test.go rename to internal/awsdevops/regions_test.go index a01546c40..feaf6c9bf 100644 --- a/internal/awsdevops/regions_token_test.go +++ b/internal/awsdevops/regions_test.go @@ -2,8 +2,6 @@ package awsdevops_test import ( "context" - "net/http" - "net/http/httptest" "testing" "github.com/stretchr/testify/assert" @@ -54,21 +52,3 @@ func TestSupportedRegionsMatchTheAWSList(t *testing.T) { "eu-central-1", "eu-west-1", "eu-west-2", "eu-west-3", "eu-north-1", }, awsdevops.SupportedRegions) } - -func TestValidateAccessTokenSendsAPIVersion(t *testing.T) { - var got []string - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - got = r.Header.Values("LD-API-Version") - if len(got) != 1 || got[0] != "20240415" { - w.WriteHeader(http.StatusBadRequest) - return - } - w.WriteHeader(http.StatusOK) - })) - t.Cleanup(server.Close) - - err := awsdevops.ValidateAccessToken(context.Background(), server.URL, "api-token") - - assert.NoError(t, err) - assert.Equal(t, []string{"20240415"}, got) -} diff --git a/internal/awsdevops/runplan_test.go b/internal/awsdevops/runplan_test.go index 1a7edc8f1..608b2f46c 100644 --- a/internal/awsdevops/runplan_test.go +++ b/internal/awsdevops/runplan_test.go @@ -4,9 +4,9 @@ import ( "context" "encoding/json" "errors" - "net/http" "strings" "testing" + "time" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/devopsagent" @@ -459,29 +459,6 @@ func TestRunPlanStopsWhenKeptAssociationIsGone(t *testing.T) { assert.NotContains(t, agent.calls, "AssociateService") } -// A plan that skips the MCP server does not register or connect one. -func TestRunPlanSkipsTheMCPServer(t *testing.T) { - agent := &fakeAgent{agentSpaces: []agenttypes.AgentSpace{namedSpace("space-7", "launchdarkly")}} - clients := newTestClients(agent, newFakeIAM()) - opts := runPlanOptions() - opts.LDAccessToken = "" - logged := collectLogs(&opts) - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - require.Equal(t, awsdevops.MCPAsk, plan.MCP.Action) - plan.MCP.Action = awsdevops.MCPSkip - - result, err := awsdevops.RunPlan(context.Background(), clients, opts, plan) - - require.NoError(t, err) - assert.Empty(t, result.MCPServiceID) - assert.NotContains(t, agent.calls, "RegisterService") - for _, in := range agent.associationInputs { - assert.Equal(t, "aws", aws.ToString(in.ServiceId)) - } - assert.Contains(t, strings.Join(*logged, "\n"), "Skipped the MCP server of the agent space") -} - // A space that reads only another account still gets this account. Setup // must not keep the association of the other account as its own. func TestRunPlanAssociatesBesideAnotherAccount(t *testing.T) { @@ -666,25 +643,19 @@ func TestConnectPlanUsesTheURLThatAWSReturns(t *testing.T) { } func TestRunPlanStopsWhenReusedMCPServerIsGone(t *testing.T) { - for name, replace := range map[string]bool{"reuse": false, "replace": true} { - t.Run(name, func(t *testing.T) { - agent := linkedSpace("mcp-old") - iamClient := newFakeIAM() - clients := newTestClients(agent, iamClient) - opts := runPlanOptions() - opts.ReplaceMCPToken = replace - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - agent.services = nil + agent := linkedSpace("mcp-old") + iamClient := newFakeIAM() + clients := newTestClients(agent, iamClient) + opts := runPlanOptions() + plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) + require.NoError(t, err) + agent.services = nil - _, err = awsdevops.RunPlan(context.Background(), clients, opts, plan) + _, err = awsdevops.RunPlan(context.Background(), clients, opts, plan) - require.ErrorIs(t, err, awsdevops.ErrPlanChanged) - assert.Empty(t, iamClient.createdRoles) - assert.NotContains(t, agent.calls, "RegisterService") - assert.NotContains(t, agent.calls, "DeregisterService") - }) - } + require.ErrorIs(t, err, awsdevops.ErrPlanChanged) + assert.Empty(t, iamClient.createdRoles) + assert.NotContains(t, agent.calls, "RegisterService") } func TestRunPlanStopsWhenMCPServerAppears(t *testing.T) { @@ -716,43 +687,6 @@ func TestRunPlanKeepsThePlannedMCPServer(t *testing.T) { assert.Equal(t, "mcp-old", result.MCPServiceID) } -// The pasted token uses the shown plan. -func TestRegisterMCPServerStopsWhenServerAppears(t *testing.T) { - agent := &fakeAgent{} - clients := newTestClients(agent, newFakeIAM()) - opts := runPlanOptions() - opts.LDAccessToken = "" - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - result, err := awsdevops.RunPlan(context.Background(), clients, opts, plan) - require.NoError(t, err) - agent.services = []agenttypes.RegisteredService{launchDarklyMCPService("mcp-someone")} - agent.associations = mcpLink("mcp-someone") - opts.LDAccessToken, opts.ReplaceMCPToken = "pasted-token", true - - err = awsdevops.RegisterMCPServer(context.Background(), clients, opts, plan.MCP, &result) - - require.ErrorIs(t, err, awsdevops.ErrPlanChanged) - assert.NotContains(t, agent.calls, "DeregisterService") - assert.NotContains(t, agent.calls, "RegisterService") -} - -func TestRegisterMCPServerRegistersAfterAsk(t *testing.T) { - agent := &fakeAgent{} - clients := newTestClients(agent, newFakeIAM()) - opts := runPlanOptions() - opts.LDAccessToken = "" - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - result, err := awsdevops.RunPlan(context.Background(), clients, opts, plan) - require.NoError(t, err) - opts.LDAccessToken = "pasted-token" - - require.NoError(t, awsdevops.RegisterMCPServer(context.Background(), clients, opts, plan.MCP, &result)) - - assert.Contains(t, agent.calls, "RegisterService") -} - func TestConnectPlanStopsWhenServerAppears(t *testing.T) { agent := &fakeAgent{agentSpaces: []agenttypes.AgentSpace{namedSpace("space-7", "platform")}} clients := newTestClients(agent, newFakeIAM()) @@ -815,127 +749,6 @@ func TestPlanRejectsBadMCPServiceID(t *testing.T) { } } -func rejectingOptions(t *testing.T, status int) awsdevops.SetupOptions { - t.Helper() - opts := runPlanOptions() - opts.LDBaseURI = stubLaunchDarkly(t, status) - - return opts -} - -// Setup checks a sent token first. -func TestRunPlanChecksTokenBeforeWrites(t *testing.T) { - for _, status := range []int{http.StatusUnauthorized, http.StatusForbidden} { - t.Run(http.StatusText(status), func(t *testing.T) { - agent, iamClient := &fakeAgent{}, newFakeIAM() - clients := newTestClients(agent, iamClient) - opts := rejectingOptions(t, status) - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - - _, err = awsdevops.RunPlan(context.Background(), clients, opts, plan) - - require.ErrorIs(t, err, awsdevops.ErrInvalidAccessToken) - assert.Empty(t, iamClient.createdRoles) - assert.NotContains(t, agent.calls, "CreateAgentSpace") - }) - } -} - -func TestRunPlanChecksReplaceTokenBeforeWrites(t *testing.T) { - agent := linkedSpace("mcp-old") - iamClient := newFakeIAM() - clients := newTestClients(agent, iamClient) - opts := rejectingOptions(t, http.StatusUnauthorized) - opts.ReplaceMCPToken = true - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - - _, err = awsdevops.RunPlan(context.Background(), clients, opts, plan) - - require.ErrorIs(t, err, awsdevops.ErrInvalidAccessToken) - assert.Empty(t, iamClient.createdRoles) -} - -// A reused server keeps its token. -func TestRunPlanSkipsTokenItWillNotSend(t *testing.T) { - agent := linkedSpace("mcp-old") - clients := newTestClients(agent, newFakeIAM()) - opts := runPlanOptions() - opts.LDBaseURI = "http://127.0.0.1:1" - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - - _, err = awsdevops.RunPlan(context.Background(), clients, opts, plan) - - require.NoError(t, err) -} - -func TestConnectPlanChecksTokenBeforeWrites(t *testing.T) { - agent := &fakeAgent{agentSpaces: []agenttypes.AgentSpace{namedSpace("space-7", "platform")}} - clients := newTestClients(agent, newFakeIAM()) - opts := rejectingOptions(t, http.StatusUnauthorized) - opts.AgentSpaceID = "space-7" - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, false) - require.NoError(t, err) - - _, err = awsdevops.ConnectPlan(context.Background(), clients, opts, plan) - - require.ErrorIs(t, err, awsdevops.ErrInvalidAccessToken) - assert.NotContains(t, agent.calls, "RegisterService") -} - -func TestRegisterMCPServerChecksPastedToken(t *testing.T) { - agent := &fakeAgent{} - clients := newTestClients(agent, newFakeIAM()) - opts := runPlanOptions() - opts.LDAccessToken = "" - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - result, err := awsdevops.RunPlan(context.Background(), clients, opts, plan) - require.NoError(t, err) - opts.LDAccessToken = "api-token" - opts.LDBaseURI = stubLaunchDarkly(t, http.StatusUnauthorized) - - err = awsdevops.RegisterMCPServer(context.Background(), clients, opts, plan.MCP, &result) - - require.ErrorIs(t, err, awsdevops.ErrInvalidAccessToken) - assert.NotContains(t, agent.calls, "RegisterService") -} - -func TestRunPlanStopsWhenReplaceSpacesChange(t *testing.T) { - agent, iamClient := replaceAcrossSpaces(), newFakeIAM() - clients := newTestClients(agent, iamClient) - opts := runPlanOptions() - opts.ReplaceMCPToken = true - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - agent.spaceLinks["space-2"] = mcpLink("mcp-old") - - _, err = awsdevops.RunPlan(context.Background(), clients, opts, plan) - - require.ErrorIs(t, err, awsdevops.ErrPlanChanged) - assert.NotContains(t, agent.calls, "DisassociateService") - assert.Empty(t, iamClient.createdRoles) -} - -// List order does not change the set. -func TestRunPlanReplaceIgnoresListOrder(t *testing.T) { - agent, iamClient := replaceAcrossSpaces(), newFakeIAM() - agent.afterListSpaces = func() { - agent.agentSpaces = []agenttypes.AgentSpace{agent.agentSpaces[2], agent.agentSpaces[1], agent.agentSpaces[0]} - } - clients := newTestClients(agent, iamClient) - opts := rejectingOptions(t, http.StatusOK) - opts.ReplaceMCPToken = true - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - - _, err = awsdevops.RunPlan(context.Background(), clients, opts, plan) - - require.NoError(t, err) -} - // The role writes can take minutes. A kept association that goes in that // time stops the run, so setup does not create a new one. func TestRunPlanStopsWhenKeptAssociationGoesDuringRoleWrites(t *testing.T) { @@ -996,91 +809,6 @@ func TestRunPlanStopsWhenReusedMCPServerGoesDuringRoleWrites(t *testing.T) { } } -// replaceRun plans and runs a token replacement across space-1 and space-3. -func replaceRun(t *testing.T, agent *fakeAgent) (awsdevops.SetupResult, error) { - t.Helper() - clients := newTestClients(agent, newFakeIAM()) - opts := rejectingOptions(t, http.StatusOK) - opts.ReplaceMCPToken = true - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - require.Equal(t, awsdevops.MCPReplace, plan.MCP.Action) - - return awsdevops.RunPlan(context.Background(), clients, opts, plan) -} - -// A failed registration on a shared server changes nothing, because the -// other spaces keep the old server and this space is still connected to it. -func TestRunPlanReplaceNameClashKeepsTheSharedServer(t *testing.T) { - agent := replaceAcrossSpaces() - agent.registerErr = errors.New("ConflictException: a service named LaunchDarkly already exists") - - _, err := replaceRun(t, agent) - - require.Error(t, err) - assert.ErrorContains(t, err, "rename it in the AWS console") - assert.NotContains(t, err.Error(), "have no LaunchDarkly MCP server") - assert.Empty(t, agent.disassociatedIDs) - assert.Empty(t, agent.deregisteredIDs) -} - -// The server of only this space is gone when the name clash happens, so the -// error names the space. -func TestRunPlanReplaceNameClashNamesTheDisconnectedSpace(t *testing.T) { - agent := replaceOnOneSpace() - agent.registerErr = errors.New("ConflictException: a service named LaunchDarkly already exists") - - _, err := replaceRun(t, agent) - - require.Error(t, err) - assert.ErrorContains(t, err, "rename it in the AWS console") - assert.ErrorContains(t, err, "agent spaces space-1 have no LaunchDarkly MCP server") - assert.ErrorContains(t, err, "--agent-space-id ") -} - -func TestRunPlanReplaceFailedDeregisterNamesTheDisconnectedSpace(t *testing.T) { - agent := replaceOnOneSpace() - agent.deregisterErr = errors.New("InternalServerException") - - _, err := replaceRun(t, agent) - - require.Error(t, err) - assert.ErrorContains(t, err, "unable to deregister the existing LaunchDarkly MCP server mcp-old") - assert.ErrorContains(t, err, "Agent spaces space-1 are already disconnected from it") - assert.ErrorContains(t, err, "--agent-space-id ") - assert.NotContains(t, agent.calls, "RegisterService") -} - -// The new server has the name of its space, which AWS needs to be unique. -func TestRunPlanReplaceNamesTheServerForTheSpace(t *testing.T) { - agent := replaceAcrossSpaces() - agent.services = []agenttypes.RegisteredService{mcpService("mcp-old", "Prod LD", awsdevops.MCPServerEndpoint)} - - _, err := replaceRun(t, agent) - - require.NoError(t, err) - require.Len(t, agent.registerInputs, 1) - details, ok := agent.registerInputs[0].ServiceDetails.(*agenttypes.ServiceDetailsMemberMcpserver) - require.True(t, ok) - assert.Equal(t, awsdevops.MCPServerNameFor("space-1"), aws.ToString(details.Value.Name)) -} - -// A link added during writes stops replace. -func TestRunPlanStopsWhenSpaceLinksDuringRoles(t *testing.T) { - agent, iamClient := replaceAcrossSpaces(), newFakeIAM() - iamClient.onCreateRole = func(string) { agent.spaceLinks["space-2"] = mcpLink("mcp-old") } - clients := newTestClients(agent, iamClient) - opts := rejectingOptions(t, http.StatusOK) - opts.ReplaceMCPToken = true - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - - _, err = awsdevops.RunPlan(context.Background(), clients, opts, plan) - - require.ErrorIs(t, err, awsdevops.ErrPlanChanged) - assert.NotContains(t, agent.calls, "DisassociateService") -} - func TestRunPlanStopsWhenAppSignInChanges(t *testing.T) { agent := &fakeAgent{ agentSpaces: []agenttypes.AgentSpace{{AgentSpaceId: aws.String("space-1"), Name: aws.String("launchdarkly")}}, @@ -1096,34 +824,6 @@ func TestRunPlanStopsWhenAppSignInChanges(t *testing.T) { assert.NotContains(t, agent.calls, "EnableOperatorApp") } -func runReplace(t *testing.T, agent *fakeAgent) error { - t.Helper() - clients := newTestClients(agent, newFakeIAM()) - opts := rejectingOptions(t, http.StatusOK) - opts.ReplaceMCPToken = true - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - - _, err = awsdevops.RunPlan(context.Background(), clients, opts, plan) - - return err -} - -// A failure on this space does not touch the other spaces. -func TestReplaceLeavesTheOtherSpacesAloneWhenThisOneFails(t *testing.T) { - agent := replaceAcrossSpaces() - agent.failAssociateFor = map[string]bool{"space-1": true} - - err := runReplace(t, agent) - - require.Error(t, err) - assert.Contains(t, err.Error(), "agent space space-1") - assert.Empty(t, agent.deregisteredIDs) - for _, in := range agent.associationInputs { - assert.NotEqual(t, "space-3", aws.ToString(in.AgentSpaceId)) - } -} - // The role step can take minutes, so the app step checks the app again. func TestRunPlanStopsWhenAppSignInChangesDuringRoles(t *testing.T) { agent := &fakeAgent{ @@ -1167,211 +867,115 @@ func linkedSpace(serviceID string) *fakeAgent { } } -// A replace on a shared server registers the new server first. If the old -// server then cannot be disconnected, setup removes the new server, so the -// space stays as it was and a run again plans the same steps. -func TestReplaceRemovesTheNewServerWhenTheDisconnectFails(t *testing.T) { - agent := replaceAcrossSpaces() - agent.disassociateErr = errors.New("InternalServerException") +// --mcp-service-id must not give a space a second LaunchDarkly server. +func TestPlanStopsWhenMCPServiceIDIsNotTheServerOfTheSpace(t *testing.T) { + agent := linkedSpace("mcp-own") + agent.services = append(agent.services, launchDarklyMCPService("mcp-chosen")) + opts := runPlanOptions() + opts.MCPServiceID = "mcp-chosen" - result, err := replaceRun(t, agent) + _, err := awsdevops.PlanSetup(context.Background(), newTestClients(agent, newFakeIAM()), opts, true) require.Error(t, err) - assert.ErrorContains(t, err, "unable to disconnect MCP server mcp-old from agent space space-1") - assert.ErrorContains(t, err, "still uses MCP server mcp-old with its old token") - assert.Equal(t, []string{"mcp-1"}, agent.deregisteredIDs) - assert.Empty(t, result.MCPServiceID) + assert.ErrorContains(t, err, "agent space space-1 already uses MCP server LaunchDarkly (mcp-own)") } -// If the new server cannot be removed either, the error names it. -func TestReplaceNamesTheNewServerItCouldNotRemove(t *testing.T) { - agent := replaceAcrossSpaces() - agent.disassociateErr = errors.New("InternalServerException") - agent.deregisterErr = errors.New("ThrottlingException") +// Setup cannot pick between two LaunchDarkly servers on one space. +func TestPlanStopsOnASpaceWithTwoLaunchDarklyServers(t *testing.T) { + agent := linkedSpace("mcp-a") + agent.services = append(agent.services, launchDarklyMCPService("mcp-b")) + agent.associations = append(agent.associations, mcpLink("mcp-b")...) - _, err := replaceRun(t, agent) + _, err := awsdevops.PlanSetup(context.Background(), newTestClients(agent, newFakeIAM()), runPlanOptions(), true) require.Error(t, err) - assert.ErrorContains(t, err, "could not remove the new MCP server mcp-1") - assert.ErrorContains(t, err, "--replace-mcp-token") -} - -// A run after such a failure removes the unused server with the name of the -// space, so the new registration does not fail on the name. -func TestReplaceRemovesAnUnusedServerWithTheNameOfTheSpace(t *testing.T) { - agent := replaceAcrossSpaces() - agent.services = append(agent.services, - mcpService("mcp-left", awsdevops.MCPServerNameFor("space-1"), awsdevops.MCPServerEndpoint)) - clients := newTestClients(agent, newFakeIAM()) - opts := rejectingOptions(t, http.StatusOK) - opts.ReplaceMCPToken = true - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - - assert.Contains(t, plan.Text().MCPServer, "Remove MCP server "+awsdevops.MCPServerNameFor("space-1")+" (mcp-left)") - _, err = awsdevops.RunPlan(context.Background(), clients, opts, plan) - require.NoError(t, err) - assert.Equal(t, []string{"mcp-left"}, agent.deregisteredIDs) - assert.Less(t, indexOf(agent.calls, "DeregisterService"), indexOf(agent.calls, "RegisterService")) -} - -// A replace gives the new server the name of the space, so a run after a -// failure can find it. The plan says so. -func TestRunPlanReplaceRenamesAnUnsharedServer(t *testing.T) { - agent := replaceOnOneSpace() - agent.services = []agenttypes.RegisteredService{mcpService("mcp-old", "Prod LD", awsdevops.MCPServerEndpoint)} - clients := newTestClients(agent, newFakeIAM()) - opts := rejectingOptions(t, http.StatusOK) - opts.ReplaceMCPToken = true - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - - assert.Equal(t, "Replace the token of MCP server Prod LD (mcp-old), and rename it to LaunchDarkly-space-1", - plan.Text().MCPServer) - _, err = awsdevops.RunPlan(context.Background(), clients, opts, plan) - require.NoError(t, err) - require.Len(t, agent.registerInputs, 1) - details, ok := agent.registerInputs[0].ServiceDetails.(*agenttypes.ServiceDetailsMemberMcpserver) - require.True(t, ok) - assert.Equal(t, awsdevops.MCPServerNameFor("space-1"), aws.ToString(details.Value.Name)) + assert.ErrorContains(t, err, "connected to 2 LaunchDarkly MCP servers") } -// A server with the name of this space that other spaces share keeps the -// name, so the new server cannot have it. The plan stops, because every run -// would fail on the name. -func TestPlanStopsWhenTheSharedServerHasTheNameOfTheSpace(t *testing.T) { - agent := replaceAcrossSpaces() - agent.services = []agenttypes.RegisteredService{ - mcpService("mcp-old", awsdevops.MCPServerNameFor("space-1"), awsdevops.MCPServerEndpoint), - } +// The plan names the other spaces that use the server. +func TestPlanWarnsWhenTheReusedServerIsShared(t *testing.T) { opts := runPlanOptions() - opts.ReplaceMCPToken = true - - _, err := awsdevops.PlanSetup(context.Background(), newTestClients(agent, newFakeIAM()), opts, true) - - require.Error(t, err) - assert.ErrorContains(t, err, "other agent spaces (space-3) use it") - assert.ErrorContains(t, err, "Rename mcp-old in the AWS console") - assert.ErrorContains(t, err, "--agent-space-id --access-token --replace-mcp-token") - assert.NotContains(t, agent.calls, "RegisterService") -} - -// A disconnect that AWS made but reported as failed is a disconnect, so the -// run keeps the new server and connects it. -func TestReplaceKeepsTheNewServerWhenTheFailedDisconnectHappened(t *testing.T) { - agent := replaceAcrossSpaces() - agent.disassociateErr = errors.New("RequestTimeout") - agent.onDisassociate = func() { delete(agent.spaceLinks, "space-1") } - result, err := replaceRun(t, agent) + plan, err := awsdevops.PlanSetup(context.Background(), newTestClients(sharedAcrossSpaces(), newFakeIAM()), opts, true) require.NoError(t, err) - assert.Empty(t, agent.deregisteredIDs) - assert.Equal(t, "mcp-1", result.MCPServiceID) - assert.Equal(t, "assoc-mcp-1", result.MCPAssociationID) -} - -// A replace on a server that only this space uses removes an unused server -// with the name of the space first, so the new server can have the name. -func TestReplaceOnAnUnsharedServerRemovesAnUnusedServerWithTheName(t *testing.T) { - agent := replaceOnOneSpace() - agent.services = append(agent.services, - mcpService("mcp-left", awsdevops.MCPServerNameFor("space-1"), awsdevops.MCPServerEndpoint)) - clients := newTestClients(agent, newFakeIAM()) - opts := rejectingOptions(t, http.StatusOK) - opts.ReplaceMCPToken = true - plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) - require.NoError(t, err) - + require.Equal(t, awsdevops.MCPReuse, plan.MCP.Action) assert.Equal(t, - "Remove MCP server "+awsdevops.MCPServerNameFor("space-1")+" (mcp-left), which no agent space uses. "+ - "Replace the token of MCP server LaunchDarkly (mcp-old), and rename it to "+awsdevops.MCPServerNameFor("space-1"), + "Reuse MCP server LaunchDarkly (mcp-old). Agent spaces space-3 also use it", plan.Text().MCPServer) - _, err = awsdevops.RunPlan(context.Background(), clients, opts, plan) - require.NoError(t, err) - assert.Equal(t, []string{"mcp-left", "mcp-old"}, agent.deregisteredIDs) } -// A server with the name of the space that another space uses stops the plan -// before setup removes the old server. -func TestReplaceOnAnUnsharedServerStopsWhenAnotherSpaceHasTheName(t *testing.T) { - agent := replaceOnOneSpace() - agent.services = append(agent.services, - mcpService("mcp-other", awsdevops.MCPServerNameFor("space-1"), awsdevops.MCPServerEndpoint)) - agent.spaceLinks["space-2"] = mcpLink("mcp-other") - opts := runPlanOptions() - opts.ReplaceMCPToken = true +func oauthMCPService(serviceID string) agenttypes.RegisteredService { + service := mcpService(serviceID, awsdevops.MCPOAuthServerName, awsdevops.MCPServerEndpoint) + details := service.AdditionalServiceDetails.(*agenttypes.AdditionalServiceDetailsMemberMcpserver) + details.Value.AuthorizationMethod = agenttypes.MCPServerAuthorizationMethodOauth3lo - _, err := awsdevops.PlanSetup(context.Background(), newTestClients(agent, newFakeIAM()), opts, true) - - require.Error(t, err) - assert.ErrorContains(t, err, "other agent spaces (space-2) use it") - assert.Empty(t, agent.deregisteredIDs) + return service } -// --mcp-service-id must not give a space a second LaunchDarkly server. -func TestPlanStopsWhenMCPServiceIDIsNotTheServerOfTheSpace(t *testing.T) { - for name, replace := range map[string]bool{"reuse": false, "replace": true} { - t.Run(name, func(t *testing.T) { - agent := linkedSpace("mcp-own") - agent.services = append(agent.services, launchDarklyMCPService("mcp-chosen")) - opts := runPlanOptions() - opts.MCPServiceID = "mcp-chosen" - opts.ReplaceMCPToken = replace - - _, err := awsdevops.PlanSetup(context.Background(), newTestClients(agent, newFakeIAM()), opts, true) +func TestOAuthRegistersAfterApproval(t *testing.T) { + defer awsdevops.SetOAuthPollInterval(time.Millisecond)() + agent := &fakeAgent{registerOutput: &devopsagent.RegisterServiceOutput{ + AdditionalStep: &agenttypes.AdditionalServiceRegistrationStepMemberOauth{ + Value: agenttypes.OAuthAdditionalStepDetails{AuthorizationUrl: aws.String("https://example.com/approve")}, + }, + }} + clients := newTestClients(agent, newFakeIAM()) + opts := runPlanOptions() + var opened string + opts.OpenURL = func(url string) error { + opened = url + // AWS registers the server after the approval. + agent.services = append(agent.services, oauthMCPService("mcp-account")) - require.Error(t, err) - assert.ErrorContains(t, err, "agent space space-1 already uses MCP server LaunchDarkly (mcp-own)") - }) + return nil } -} - -// Setup cannot pick between two LaunchDarkly servers on one space. -func TestPlanStopsOnASpaceWithTwoLaunchDarklyServers(t *testing.T) { - agent := linkedSpace("mcp-a") - agent.services = append(agent.services, launchDarklyMCPService("mcp-b")) - agent.associations = append(agent.associations, mcpLink("mcp-b")...) + plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) + require.NoError(t, err) + require.Equal(t, awsdevops.MCPRegister, plan.MCP.Action) - _, err := awsdevops.PlanSetup(context.Background(), newTestClients(agent, newFakeIAM()), runPlanOptions(), true) + result, err := awsdevops.RunPlan(context.Background(), clients, opts, plan) - require.Error(t, err) - assert.ErrorContains(t, err, "connected to 2 LaunchDarkly MCP servers") + require.NoError(t, err) + assert.Equal(t, "https://example.com/approve", opened) + assert.Equal(t, "mcp-account", result.MCPServiceID) + require.Len(t, agent.registerInputs, 1) + details := agent.registerInputs[0].ServiceDetails.(*agenttypes.ServiceDetailsMemberMcpserver).Value + assert.Equal(t, awsdevops.MCPOAuthServerName, aws.ToString(details.Name)) + assert.IsType(t, &agenttypes.MCPServerAuthorizationConfigMemberAuthorizationDiscovery{}, details.AuthorizationConfig) + assert.Contains(t, agent.calls, "AssociateService") } -// A space that shares a server with other spaces acts as the account of its -// token, so the plan says so. -func TestPlanWarnsWhenTheReusedServerIsShared(t *testing.T) { +func TestOAuthReusesTheServerOfTheAccount(t *testing.T) { + agent := &fakeAgent{ + agentSpaces: []agenttypes.AgentSpace{namedSpace("space-7", "platform")}, + services: []agenttypes.RegisteredService{oauthMCPService("mcp-account")}, + } + clients := newTestClients(agent, newFakeIAM()) opts := runPlanOptions() - opts.LDAccessToken = "" - - plan, err := awsdevops.PlanSetup(context.Background(), newTestClients(replaceAcrossSpaces(), newFakeIAM()), opts, true) - + opts.AgentSpaceID = "space-7" + plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, false) require.NoError(t, err) require.Equal(t, awsdevops.MCPReuse, plan.MCP.Action) - assert.Equal(t, - "Reuse MCP server LaunchDarkly (mcp-old). Agent spaces space-3 also use it and its token. "+ - "To give this agent space its own server, pass --access-token and --replace-mcp-token", - plan.Text().MCPServer) -} -// AWS allows only letters, digits, "_" and "-" in an MCP server name. -func TestMCPServerNameForHasOnlyAllowedCharacters(t *testing.T) { - name := awsdevops.MCPServerNameFor("0b9f6c3e-5d7a-4c1e-9a2b-3f4e5d6c7b8a") + result, err := awsdevops.ConnectPlan(context.Background(), clients, opts, plan) - assert.Regexp(t, `^[a-zA-Z0-9_-]{1,128}$`, name) + require.NoError(t, err) + assert.Equal(t, "mcp-account", result.MCPServiceID) + assert.NotContains(t, agent.calls, "RegisterService") } -// The name keeps the last group of a UUID space ID. -func TestMCPServerNameForKeepsTheEndOfTheSpaceID(t *testing.T) { - assert.Equal(t, "LaunchDarkly-3f4e5d6c7b8a", awsdevops.MCPServerNameFor("0b9f6c3e-5d7a-4c1e-9a2b-3f4e5d6c7b8a")) -} +// A server with a token is never the OAuth server of the account, even with +// the name that the console guide uses. +func TestOAuthDoesNotReuseAServerWithAToken(t *testing.T) { + token := launchDarklyMCPService("mcp-token") + details := token.AdditionalServiceDetails.(*agenttypes.AdditionalServiceDetailsMemberMcpserver) + details.Value.AuthorizationMethod = agenttypes.MCPServerAuthorizationMethodBearerToken + agent := &fakeAgent{services: []agenttypes.RegisteredService{token}} + clients := newTestClients(agent, newFakeIAM()) -// AWS caps the server name, "_" and a tool name at 64 characters. -func TestMCPServerNameForLeavesRoomForEveryTool(t *testing.T) { - name := awsdevops.MCPServerNameFor("0b9f6c3e-5d7a-4c1e-9a2b-3f4e5d6c7b8a") + plan, err := awsdevops.PlanSetup(context.Background(), clients, runPlanOptions(), true) - for _, tool := range awsdevops.MCPTools() { - assert.LessOrEqual(t, len(name+"_"+tool), 64, tool) - } + require.NoError(t, err) + assert.Equal(t, awsdevops.MCPRegister, plan.MCP.Action) } diff --git a/internal/awsdevops/setup.go b/internal/awsdevops/setup.go index e3deb7e16..ddc0f10bf 100644 --- a/internal/awsdevops/setup.go +++ b/internal/awsdevops/setup.go @@ -34,8 +34,6 @@ const ( // AWS account association. awsServiceID = "aws" - DefaultLDBaseURI = "https://app.launchdarkly.com" - roleAssumableTimeout = 2 * time.Minute roleAssumableFirstWait = time.Second roleAssumableMaxInterval = 10 * time.Second @@ -43,12 +41,26 @@ const ( MCPServerName = "LaunchDarkly" MCPServerEndpoint = "https://mcp.launchdarkly.com/mcp/launchdarkly" + // MCPOAuthServerName is the name of the MCP server that setup registers. + // It differs from MCPServerName, the name that the AWS console guide + // uses, so it does not conflict with a server that someone registered + // with a token. + MCPOAuthServerName = "LaunchDarkly-OAuth" + + // oauthTimeout is how long setup waits for a person to approve access. + oauthTimeout = 10 * time.Minute + agentSpaceDescription = "Managed by the LaunchDarkly CLI" + mcpServerDescription = "LaunchDarkly feature flag management MCP server" // ElevatedActionsPreference lets the agent run actions that change AWS // resources. ElevatedActionsPreference = "elevatedActionsEnabled" ) +// oauthPollInterval is how often setup looks for the server that AWS +// registers after a person approves access. Tests make it shorter. +var oauthPollInterval = 5 * time.Second + // The agent invokes read-only tools without asking for approval, so anything // that changes LaunchDarkly state has to be classified explicitly. The lists // cover the steps of the experiment-orchestration skill: find a flag, create @@ -109,17 +121,15 @@ type SetupOptions struct { IdpClientID string IdpClientSecret string - LDBaseURI string - LDAccessToken string - ReplaceMCPToken bool - - // MCPServiceID, when set, is the registration to reuse or replace instead - // of the server of the agent space. + // MCPServiceID, when set, is the registration to connect instead of the + // server of the AWS account. MCPServiceID string // MCPEndpoint overrides the LaunchDarkly MCP endpoint, for an account // running the open-source MCP server itself. The endpoint identifies the // registration, so a self-hosted one is never confused with ours. MCPEndpoint string + // OpenURL, when set, opens the OAuth approval page in a browser. + OpenURL func(url string) error // Logf, when set, reports progress as each step completes. Logf func(format string, args ...any) @@ -349,10 +359,6 @@ type MCPAction int const ( MCPRegister MCPAction = iota MCPReuse - MCPReplace - MCPAsk - // MCPSkip means that the user gave no token, so the run skips the server. - MCPSkip ) // MCPStep is the planned MCP server step. @@ -361,17 +367,12 @@ type MCPStep struct { ServiceID string // Endpoint is the URL that the server serves. Endpoint string - // Name is the name of an existing server in AWS. The customer can change - // it, so it can differ from MCPServerNameFor. + // Name is the name of an existing server in AWS. Name string // Spaces lists spaces linked to the server. Spaces []string // SpaceID is the agent space of the step. It is empty for a new space. SpaceID string - // Unused is a server with the name that setup gives the space, which no - // space uses. A replace on a shared server removes it first, so that the - // new server can have the name. An earlier run that failed leaves one. - Unused string } // shared reports whether other spaces use the server. @@ -380,7 +381,7 @@ func (m MCPStep) shared() bool { } // otherSpaces lists the spaces, other than the space of the step, that use -// the server. A token replacement does not change them. +// the server. func (m MCPStep) otherSpaces() []string { var others []string for _, id := range m.Spaces { @@ -672,7 +673,7 @@ func planMCPStep(ctx context.Context, clients Clients, opts SetupOptions, spaceI } func planMCPAction(ctx context.Context, clients Clients, opts SetupOptions, spaceID string) (MCPStep, error) { - linked, own, err := spaceMCPServers(ctx, clients, spaceID, opts.mcpEndpoint()) + linked, err := linkedMCPServer(ctx, clients, spaceID, opts.mcpEndpoint()) if err != nil { return MCPStep{}, err } @@ -694,62 +695,24 @@ func planMCPAction(ctx context.Context, clients Clients, opts SetupOptions, spac } existing = opts.MCPServiceID case existing == "": - // A run that registered the server and then failed to connect it - // leaves a server with the name of the space. - existing, name = own.ID, own.Name - } - - switch { - case existing != "" && opts.ReplaceMCPToken && opts.LDAccessToken != "": - spaces, err := spacesWith(ctx, clients, existing) + // Every agent space uses the one server of the account. + account, err := accountMCPServer(ctx, clients, opts.mcpEndpoint()) if err != nil { return MCPStep{}, err } - step := MCPStep{Action: MCPReplace, ServiceID: existing, Name: name, Spaces: spaces, SpaceID: spaceID} - // The new server gets the name of the space. A server that other - // spaces use keeps the name, so the new server cannot have it. - if step.shared() && own.ID == existing { - return MCPStep{}, nameTaken(own, spaceID, step.otherSpaces()) - } - if own.ID != "" && own.ID != existing { - users, err := spacesWith(ctx, clients, own.ID) - if err != nil { - return MCPStep{}, err - } - if len(users) > 0 { - return MCPStep{}, nameTaken(own, spaceID, users) - } - // No space uses it, so the run removes it before it registers. - step.Unused = own.ID - } - - return step, nil - case existing != "": - // The other spaces are only for the plan text, which warns that the - // space shares a token. - spaces, err := spacesWith(ctx, clients, existing) - if err != nil { - return MCPStep{}, err - } - - return MCPStep{Action: MCPReuse, ServiceID: existing, Name: name, Spaces: spaces}, nil - case opts.LDAccessToken != "": + existing, name = account.ID, account.Name + } + if existing == "" { return MCPStep{Action: MCPRegister}, nil - default: - return MCPStep{Action: MCPAsk}, nil } -} -// nameTaken is the error for a server that has the name that setup gives an -// agent space, while other agent spaces use it. -func nameTaken(own mcpRef, spaceID string, users []string) error { - return fmt.Errorf( - "MCP server %s (%s) has the name that setup gives agent space %s, and other agent spaces (%s) use it, "+ - "so setup cannot register a new server with that name. Rename %s in the AWS console, then run setup again. "+ - "Or, to move those agent spaces to their own servers, run setup with --agent-space-id "+ - "--access-token --replace-mcp-token for each one", - own.Name, own.ID, spaceID, strings.Join(users, ", "), own.ID, - ) + // The other spaces are only for the plan text. + spaces, err := spacesWith(ctx, clients, existing) + if err != nil { + return MCPStep{}, err + } + + return MCPStep{Action: MCPReuse, ServiceID: existing, Name: name, Spaces: spaces}, nil } // checkMCPServiceID needs a registered MCP server at the LaunchDarkly @@ -781,7 +744,7 @@ func checkMCPServiceID(ctx context.Context, clients Clients, id, endpoint string } // name is the name of the existing server in AWS, or MCPServerName when AWS -// sent none. A server that setup registers gets MCPServerNameFor instead. +// sent none. func (m MCPStep) name() string { if m.Name != "" { return m.Name @@ -791,99 +754,51 @@ func (m MCPStep) name() string { } func (m MCPStep) describe() string { - switch m.Action { - case MCPReplace: - var line string - if m.shared() { - line = fmt.Sprintf( - "Register a new MCP server for this agent space at %s. MCP server %s (%s) keeps its token for %s", - m.Endpoint, m.name(), m.ServiceID, strings.Join(m.otherSpaces(), ", "), - ) - } else { - line = fmt.Sprintf("Replace the token of MCP server %s (%s)", m.name(), m.ServiceID) - // Every server that setup registers has the name of its space, - // so a run after a failure can find it. - switch { - case m.SpaceID == "": - line += ", and give it the name of the new agent space" - case m.name() != MCPServerNameFor(m.SpaceID): - line += ", and rename it to " + MCPServerNameFor(m.SpaceID) - } - } - if m.Unused != "" { - line = fmt.Sprintf("Remove MCP server %s (%s), which no agent space uses. %s", - MCPServerNameFor(m.SpaceID), m.Unused, line) - } - - return line - case MCPReuse: - line := fmt.Sprintf("Reuse MCP server %s (%s)", m.name(), m.ServiceID) - if m.shared() { - line += fmt.Sprintf( - ". Agent spaces %s also use it and its token. "+ - "To give this agent space its own server, pass --access-token and --replace-mcp-token", - strings.Join(m.otherSpaces(), ", "), - ) - } - - return line - case MCPRegister: - return fmt.Sprintf("Register a new MCP server for this agent space at %s", m.Endpoint) - case MCPSkip: - return "Skip the MCP server of this agent space, because no token was given" - default: - return "Ask for a token, then register a new MCP server for this agent space" + if m.Action == MCPRegister { + return fmt.Sprintf( + "Register MCP server %s for the AWS account at %s. You approve access to LaunchDarkly in a browser", + MCPOAuthServerName, m.Endpoint, + ) } -} - -// AWS stores tokens write-only. Check them first. -func checkSentToken(ctx context.Context, opts SetupOptions, step MCPStep) error { - if step.Action != MCPRegister && step.Action != MCPReplace { - return nil + line := fmt.Sprintf("Reuse MCP server %s (%s)", m.name(), m.ServiceID) + if m.shared() { + line += fmt.Sprintf(". Agent spaces %s also use it", strings.Join(m.otherSpaces(), ", ")) } - return CheckToken(ctx, opts.LDBaseURI, opts.LDAccessToken) + return line } // checkMCP compares AWS with the MCP step of agent space spaceID. The ID is // empty for a space that the run creates. func checkMCP(ctx context.Context, clients Clients, opts SetupOptions, step MCPStep, spaceID string) error { - if step.Action == MCPSkip { - return nil - } - if step.Action == MCPReuse || step.Action == MCPReplace { + if step.Action == MCPReuse { services, err := listServices(ctx, clients) if err != nil { return err } - found := false for _, service := range services { - found = found || aws.ToString(service.ServiceId) == step.ServiceID - } - if !found { - return planChanged("MCP server %s disappeared after the plan", step.ServiceID) - } - if step.Action != MCPReplace { - return nil - } - now, err := spacesWith(ctx, clients, step.ServiceID) - if err != nil { - return err - } - if !slices.Equal(now, step.Spaces) { - return planChanged("the spaces using MCP server %s changed", step.ServiceID) + if aws.ToString(service.ServiceId) == step.ServiceID { + return nil + } } - return nil + return planChanged("MCP server %s disappeared after the plan", step.ServiceID) } - linked, own, err := spaceMCPServers(ctx, clients, spaceID, opts.mcpEndpoint()) + linked, err := linkedMCPServer(ctx, clients, spaceID, opts.mcpEndpoint()) if err != nil { return err } - if linked.ID != "" || own.ID != "" { + if linked.ID != "" { return planChanged("agent space %s got an MCP server at %s after the plan", spaceID, opts.mcpEndpoint()) } + account, err := accountMCPServer(ctx, clients, opts.mcpEndpoint()) + if err != nil { + return err + } + if account.ID != "" { + return planChanged("the AWS account got MCP server %s (%s) after the plan", account.Name, account.ID) + } return nil } @@ -923,9 +838,6 @@ func RunPlan(ctx context.Context, clients Clients, opts SetupOptions, plan Setup if err := checkMCP(ctx, clients, opts, plan.MCP, plan.Space.ID); err != nil { return result, err } - if err := checkSentToken(ctx, opts, plan.MCP); err != nil { - return result, err - } spaceRole, err := checkRole(ctx, clients.IAM, plan, AgentSpaceRoleName) if err != nil { return result, err @@ -1044,23 +956,6 @@ func (o SetupOptions) mcpEndpoint() string { return MCPServerEndpoint } -// registrationFailed names the agent spaces a failed token replacement left -// without an MCP server: the old registration is already gone by then, and its -// token cannot be read back, so only a re-run with a working token restores them. -func registrationFailed(err error, detached []string) error { - if len(detached) == 0 { - return fmt.Errorf("unable to register the LaunchDarkly MCP server: %w", err) - } - - return fmt.Errorf( - "unable to register the LaunchDarkly MCP server: %w. The previous registration was already "+ - "removed, so agent spaces %s have no LaunchDarkly MCP server. To connect each one again, "+ - "run setup with --access-token --agent-space-id ", - err, - strings.Join(detached, ", "), - ) -} - // ConnectPlan runs a connect plan. func ConnectPlan(ctx context.Context, clients Clients, opts SetupOptions, plan SetupPlan) (SetupResult, error) { logf := opts.Logf @@ -1103,9 +998,6 @@ func ConnectPlan(ctx context.Context, clients Clients, opts SetupOptions, plan S if err := checkMCP(ctx, clients, opts, plan.MCP, plan.Space.ID); err != nil { return result, err } - if err := checkSentToken(ctx, opts, plan.MCP); err != nil { - return result, err - } if err := registerMCPServer(ctx, clients, opts, plan.MCP, &result, logf); err != nil { return result, err } @@ -1197,26 +1089,6 @@ func isRoleNotAssumableYet(err error) bool { return strings.Contains(err.Error(), "Invalid STS role configuration") } -// RegisterMCPServer registers and associates the LaunchDarkly MCP server on -// its own, for a token collected after Setup has run. -func RegisterMCPServer(ctx context.Context, clients Clients, opts SetupOptions, step MCPStep, result *SetupResult) error { - logf := opts.Logf - if logf == nil { - logf = func(string, ...any) {} - } - if err := checkMCP(ctx, clients, opts, step, result.AgentSpaceID); err != nil { - return err - } - if step.Action == MCPAsk { - step.Action = MCPRegister - } - if err := checkSentToken(ctx, opts, step); err != nil { - return err - } - - return registerMCPServer(ctx, clients, opts, step, result, logf) -} - func registerMCPServer( ctx context.Context, clients Clients, @@ -1225,188 +1097,97 @@ func registerMCPServer( result *SetupResult, logf func(string, ...any), ) error { - endpoint := opts.mcpEndpoint() - existing := step.ServiceID - switch step.Action { - case MCPAsk: - // The caller asks for a token. - return nil - case MCPSkip: - logf("Skipped the MCP server of the agent space, because no token was given") - - return nil - } if step.Action == MCPReuse { - result.MCPServiceID = existing - logf("Reused MCP server %s (%s) with its current token", step.name(), existing) - logf("Pass --access-token and --replace-mcp-token to change the token") + result.MCPServiceID = step.ServiceID + logf("Reused MCP server %s (%s)", step.name(), step.ServiceID) return associateMCPServer(ctx, clients, step.name(), result, logf) } - // A run after a failure finds an unconnected server by this name. - name := MCPServerNameFor(result.AgentSpaceID) - var detached []string - if existing != "" { - // Narrow the window after other writes. - now, err := spacesWith(ctx, clients, existing) - if err != nil { - return err - } - if !slices.Equal(now, step.Spaces) { - return planChanged("the spaces using MCP server %s changed", existing) - } - } - if step.Unused != "" { - if err := removeUnused(ctx, clients, step.Unused, logf); err != nil { - return err - } - } - // Other spaces keep a shared server with its token. Only a server that no - // other space uses is removed. - if existing != "" && !step.shared() { - var err error - detached, err = disassociateEverywhere(ctx, clients, existing, logf) - if err != nil { - return err - } - if _, err := clients.Agent.DeregisterService(ctx, &devopsagent.DeregisterServiceInput{ - ServiceId: aws.String(existing), - }); err != nil { - err = fmt.Errorf("unable to deregister the existing LaunchDarkly MCP server %s: %w", existing, err) - if len(detached) > 0 { - // The server is still there with its old token. - err = fmt.Errorf( - "%w. Agent spaces %s are already disconnected from it. "+ - "To connect each one again, run setup with --agent-space-id ", - err, strings.Join(detached, ", "), - ) - } - return err - } - logf("Removed MCP server %s (%s) to replace its token", step.name(), existing) - } + return registerOAuthMCPServer(ctx, clients, opts, result, logf) +} +// registerOAuthMCPServer registers the MCP server of the AWS account with +// OAuth. AWS does not register the server at once. It sends an approval URL, +// and it registers the server after a person approves access in a browser. +// Setup then finds the server by its endpoint and its OAuth sign-in. +func registerOAuthMCPServer( + ctx context.Context, + clients Clients, + opts SetupOptions, + result *SetupResult, + logf func(string, ...any), +) error { + endpoint := opts.mcpEndpoint() input := &devopsagent.RegisterServiceInput{ Service: agenttypes.PostRegisterServiceSupportedServiceMcpServer, ServiceDetails: &agenttypes.ServiceDetailsMemberMcpserver{ Value: agenttypes.MCPServerDetails{ - Name: aws.String(name), + Name: aws.String(MCPOAuthServerName), Endpoint: aws.String(endpoint), - Description: aws.String("LaunchDarkly feature flag management MCP server"), - AuthorizationConfig: &agenttypes.MCPServerAuthorizationConfigMemberBearerToken{ - Value: agenttypes.MCPServerBearerTokenConfig{ - TokenName: aws.String("launchdarkly-api-token"), - TokenValue: aws.String(opts.LDAccessToken), + Description: aws.String(mcpServerDescription), + AuthorizationConfig: &agenttypes.MCPServerAuthorizationConfigMemberAuthorizationDiscovery{ + Value: agenttypes.MCPServerAuthorizationDiscoveryConfig{ + ReturnToEndpoint: aws.String(ConsoleURL(clients.Region)), }, }, }, }, } - registration, err := registerWithRetry(ctx, clients, input, logf, endpoint) if err != nil { - if strings.Contains(err.Error(), "already exists") { - // The plan removes or stops on a server at this endpoint with the - // name, so the server is at another endpoint, or it appeared - // after the plan. - err = fmt.Errorf( - "an MCP server named %s is already registered on this account, so setup cannot register "+ - "one at %s with that name. If that server is at another endpoint, rename it in the AWS console. "+ - "Otherwise, run setup again: %w", - name, endpoint, err, - ) - } - - return registrationFailed(err, detached) + return fmt.Errorf("unable to register the LaunchDarkly MCP server: %w", err) } result.MCPServiceID = aws.ToString(registration.ServiceId) - logf("Registered MCP server %s (%s)", name, result.MCPServiceID) - - if existing != "" && step.shared() { - if err := leaveSharedServer(ctx, clients, step, result, logf); err != nil { + if result.MCPServiceID == "" { + step, ok := registration.AdditionalStep.(*agenttypes.AdditionalServiceRegistrationStepMemberOauth) + if !ok { + return errors.New("AWS did not register the LaunchDarkly MCP server and sent no OAuth approval URL") + } + approvalURL := aws.ToString(step.Value.AuthorizationUrl) + logf("To connect the MCP server, approve access to LaunchDarkly at:\n %s", approvalURL) + if opts.OpenURL != nil { + if err := opts.OpenURL(approvalURL); err != nil { + logf("Unable to open a browser. Open the URL yourself: %v", err) + } + } + result.MCPServiceID, err = waitForOAuthServer(ctx, clients, endpoint, logf) + if err != nil { return err } } + logf("Registered MCP server %s (%s)", MCPOAuthServerName, result.MCPServiceID) - return associateMCPServer(ctx, clients, name, result, logf) + return associateMCPServer(ctx, clients, MCPOAuthServerName, result, logf) } -// leaveSharedServer disconnects a shared server from this space only, after -// the new server is registered. The other spaces keep it with its token. If -// the disconnect fails, it removes the new server, so that the space stays as -// it was and a run again plans the same steps. -func leaveSharedServer(ctx context.Context, clients Clients, step MCPStep, result *SetupResult, logf func(string, ...any)) error { - link, err := FindAssociation(ctx, clients, result.AgentSpaceID, step.ServiceID) - if err == nil && link != "" { - _, err = clients.Agent.DisassociateService(ctx, &devopsagent.DisassociateServiceInput{ - AgentSpaceId: aws.String(result.AgentSpaceID), - AssociationId: aws.String(link), - }) - } - var checkErr error - if err != nil { - // A disconnect can fail in the reply after AWS made it, so look - // again before the new server is removed. - var still string - still, checkErr = FindAssociation(ctx, clients, result.AgentSpaceID, step.ServiceID) - if checkErr == nil && still == "" { - err = nil +// waitForOAuthServer looks for the OAuth server of the account until it +// appears or oauthTimeout passes. +func waitForOAuthServer(ctx context.Context, clients Clients, endpoint string, logf func(string, ...any)) (string, error) { + logf("Waiting for you to approve access in the browser") + deadline := time.Now().Add(oauthTimeout) + for { + server, err := accountMCPServer(ctx, clients, endpoint) + if err != nil { + return "", err } - } - if err == nil { - if link != "" { - logf("Disconnected MCP server %s (%s) from the agent space. Agent spaces %s keep it", - step.name(), step.ServiceID, strings.Join(step.otherSpaces(), ", ")) + if server.ID != "" { + return server.ID, nil + } + if time.Now().After(deadline) { + return "", fmt.Errorf( + "AWS did not register the LaunchDarkly MCP server in %s. If you did not approve access, run setup again", + oauthTimeout, + ) } - return nil - } - - err = fmt.Errorf("unable to disconnect MCP server %s from agent space %s: %w", step.ServiceID, result.AgentSpaceID, err) - created := result.MCPServiceID - result.MCPServiceID = "" - if _, removeErr := clients.Agent.DeregisterService(ctx, &devopsagent.DeregisterServiceInput{ - ServiceId: aws.String(created), - }); removeErr != nil { - // A run again removes the new server, because no space uses it. - return fmt.Errorf( - "%w. Setup also could not remove the new MCP server %s, which no agent space uses. "+ - "Run setup again with --replace-mcp-token: %w", - err, created, removeErr, - ) - } - if checkErr != nil { - logf("Removed the new MCP server (%s)", created) - - return fmt.Errorf("%w. Setup could not check if the agent space still uses MCP server %s: %w", - err, step.ServiceID, checkErr) - } - - logf("Removed the new MCP server (%s), because the agent space still uses MCP server %s", created, step.ServiceID) - - return fmt.Errorf("%w. The agent space still uses MCP server %s with its old token", err, step.ServiceID) -} - -// removeUnused removes a server that no space uses, after a check that this -// is still true. -func removeUnused(ctx context.Context, clients Clients, serviceID string, logf func(string, ...any)) error { - users, err := spacesWith(ctx, clients, serviceID) - if err != nil { - return err - } - if len(users) > 0 { - return planChanged("agent spaces %s started to use MCP server %s after the plan", strings.Join(users, ", "), serviceID) - } - if _, err := clients.Agent.DeregisterService(ctx, &devopsagent.DeregisterServiceInput{ - ServiceId: aws.String(serviceID), - }); err != nil { - return fmt.Errorf("unable to remove MCP server %s, which no agent space uses: %w", serviceID, err) + select { + case <-ctx.Done(): + return "", ctx.Err() + case <-time.After(oauthPollInterval): + } } - logf("Removed MCP server (%s), which no agent space used", serviceID) - - return nil } // spacesWith lists spaces associated with serviceID. @@ -1431,54 +1212,6 @@ func spacesWith(ctx context.Context, clients Clients, serviceID string) ([]strin return ids, nil } -// disassociateEverywhere removes a service from every agent space that uses -// it, which AWS requires before the service can be deregistered. It returns -// the agent spaces it was removed from, so the re-registered service can be -// reconnected to them. -func disassociateEverywhere( - ctx context.Context, - clients Clients, - serviceID string, - logf func(string, ...any), -) ([]string, error) { - spaces, err := listAgentSpaces(ctx, clients) - if err != nil { - return nil, err - } - - var detached []string - for _, space := range spaces { - agentSpaceID := aws.ToString(space.AgentSpaceId) - associationID, err := FindAssociation(ctx, clients, agentSpaceID, serviceID) - if err != nil { - return nil, err - } - if associationID == "" { - continue - } - if _, err := clients.Agent.DisassociateService(ctx, &devopsagent.DisassociateServiceInput{ - AgentSpaceId: aws.String(agentSpaceID), - AssociationId: aws.String(associationID), - }); err != nil { - err = fmt.Errorf("unable to disassociate service %s from agent space %s: %w", serviceID, agentSpaceID, err) - if len(detached) > 0 { - err = fmt.Errorf( - "%w. Agent spaces %s are already disconnected from it. "+ - "To connect each one again, run setup with --agent-space-id ", - err, strings.Join(detached, ", "), - ) - } - - return nil, err - } - logf("Disassociated service (%s) from agent space (%s)", serviceID, agentSpaceID) - detached = append(detached, agentSpaceID) - } - - return detached, nil -} - -// associateMCPServer connects the server, named name in AWS, to the space. func associateMCPServer( ctx context.Context, clients Clients, @@ -1731,15 +1464,6 @@ func ConsoleURL(region string) string { return fmt.Sprintf("https://%s.console.aws.amazon.com/aidevops/home?region=%[1]s", region) } -// AccessTokenURL is the LaunchDarkly page where a service token is created. -func AccessTokenURL(baseURI string) string { - if baseURI == "" { - baseURI = DefaultLDBaseURI - } - - return strings.TrimSuffix(baseURI, "/") + "/settings/authorization" -} - // OperatorAppURL is the browser entry point for an agent space's operator app. func OperatorAppURL(agentSpaceID string) string { return fmt.Sprintf("https://%s.aidevops.global.app.aws", agentSpaceID) diff --git a/internal/awsdevops/setup_test.go b/internal/awsdevops/setup_test.go index d2fadc469..2c5c28301 100644 --- a/internal/awsdevops/setup_test.go +++ b/internal/awsdevops/setup_test.go @@ -3,7 +3,6 @@ package awsdevops_test import ( "context" "fmt" - "net/http" "strings" "testing" @@ -37,29 +36,17 @@ func TestSetupProgressLinesHaveNoARNs(t *testing.T) { "new space with a slow role": { agent: &fakeAgent{associateFailures: 1}, iam: newFakeIAM(), - opts: awsdevops.SetupOptions{AgentSpaceName: "launchdarkly", AuthFlow: "iam", LDAccessToken: "api-token"}, + opts: awsdevops.SetupOptions{AgentSpaceName: "launchdarkly", AuthFlow: "iam"}, }, "roles ldcli did not create": { agent: &fakeAgent{}, iam: unmanaged, opts: awsdevops.SetupOptions{AgentSpaceID: "space-1", AuthFlow: "iam"}, }, - "token replacement": { - agent: &fakeAgent{ - services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-old")}, - associations: []agenttypes.Association{{ - AssociationId: aws.String("assoc-old"), - ServiceId: aws.String("mcp-old"), - }}, - }, - iam: newFakeIAM(), - opts: awsdevops.SetupOptions{ - AgentSpaceID: "space-1", - AuthFlow: "iam", - LDAccessToken: "api-token", - ReplaceMCPToken: true, - LDBaseURI: stubLaunchDarkly(t, http.StatusOK), - }, + "reused account server": { + agent: &fakeAgent{services: []agenttypes.RegisteredService{oauthMCPService("mcp-account")}}, + iam: newFakeIAM(), + opts: awsdevops.SetupOptions{AgentSpaceID: "space-1", AuthFlow: "iam"}, }, } for name, s := range scenarios { @@ -79,7 +66,7 @@ func TestSetupProgressLinesHaveNoARNs(t *testing.T) { } func TestConnectExistingProgressHasNoARNs(t *testing.T) { - opts := awsdevops.SetupOptions{AgentSpaceID: "space-2", LDAccessToken: "api-token"} + opts := awsdevops.SetupOptions{AgentSpaceID: "space-2"} logged := collectLogs(&opts) agent := &fakeAgent{services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")}} @@ -113,7 +100,6 @@ func TestPlanSetupMakesNoWrites(t *testing.T) { _, err := awsdevops.PlanSetup(context.Background(), newTestClients(agent, iamClient), awsdevops.SetupOptions{ AgentSpaceName: "launchdarkly", - LDAccessToken: "api-token", }, true) require.NoError(t, err) @@ -199,27 +185,16 @@ func TestPlanSetupNamesTheMCPServerAction(t *testing.T) { opts awsdevops.SetupOptions want []string }{ - "reuse without a token": { + "reuse the server of the space": { services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")}, want: []string{"reuse", "mcp-existing"}, }, - "reuse when replace has no token": { - services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")}, - opts: awsdevops.SetupOptions{ReplaceMCPToken: true}, - want: []string{"reuse", "mcp-existing"}, - }, - "replace the token": { - services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")}, - opts: awsdevops.SetupOptions{ReplaceMCPToken: true, LDAccessToken: "api-token"}, - want: []string{"replace", "mcp-existing"}, - }, - "register with a token": { - opts: awsdevops.SetupOptions{LDAccessToken: "api-token"}, - want: []string{"register", awsdevops.MCPServerEndpoint}, + "register with OAuth": { + want: []string{"register", awsdevops.MCPOAuthServerName, awsdevops.MCPServerEndpoint, "browser"}, }, "register at another endpoint": { services: []agenttypes.RegisteredService{launchDarklyMCPService("mcp-existing")}, - opts: awsdevops.SetupOptions{LDAccessToken: "api-token", MCPEndpoint: "https://mcp.example.com/mcp"}, + opts: awsdevops.SetupOptions{MCPEndpoint: "https://mcp.example.com/mcp"}, want: []string{"register", "https://mcp.example.com/mcp"}, }, } diff --git a/internal/awsdevops/token.go b/internal/awsdevops/token.go deleted file mode 100644 index fcf915d8f..000000000 --- a/internal/awsdevops/token.go +++ /dev/null @@ -1,63 +0,0 @@ -package awsdevops - -import ( - "context" - "errors" - "fmt" - "io" - "net/http" - "net/url" - "time" - - "github.com/launchdarkly/ldcli/internal/resources" -) - -// ErrInvalidAccessToken is the sentinel for a token LaunchDarkly rejects. -var ErrInvalidAccessToken = errors.New("LaunchDarkly rejected the access token") - -var ldHTTPClient = &http.Client{Timeout: 15 * time.Second} - -// CheckToken checks a token. Tests replace it. -var CheckToken = ValidateAccessToken - -// ValidateAccessToken checks a token against the LaunchDarkly API before AWS -// stores it. AWS keeps the token write-only, so replacing a working -// registration with an unusable token cannot be undone. -func ValidateAccessToken(ctx context.Context, baseURI, token string) error { - if baseURI == "" { - baseURI = DefaultLDBaseURI - } - endpoint, err := url.JoinPath(baseURI, "api/v2/caller-identity") - if err != nil { - return fmt.Errorf("invalid LaunchDarkly base URI %q: %w", baseURI, err) - } - - req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) - if err != nil { - return err - } - req.Header.Set("Authorization", token) - req.Header.Set("LD-API-Version", resources.APIVersion) - - res, err := ldHTTPClient.Do(req) - if err != nil { - return fmt.Errorf("unable to reach LaunchDarkly at %s to check the access token: %w", baseURI, err) - } - defer func() { - _, _ = io.Copy(io.Discard, res.Body) - _ = res.Body.Close() - }() - - switch { - case res.StatusCode == http.StatusUnauthorized || res.StatusCode == http.StatusForbidden: - return fmt.Errorf( - "%w. Create a service token at %s and pass it with --access-token", - ErrInvalidAccessToken, - AccessTokenURL(baseURI), - ) - case res.StatusCode >= http.StatusBadRequest: - return fmt.Errorf("unable to check the access token: LaunchDarkly returned %s", res.Status) - } - - return nil -} From 9db7e6049e78e5e86a0c0d4c691eee35a14cfc14 Mon Sep 17 00:00:00 2001 From: Christina Chu Date: Wed, 28 Oct 2026 12:59:33 -0700 Subject: [PATCH 2/6] test: cover the poll loop of the OAuth registration The new test adds the server on the third poll, so setup must wait and poll again. The approval test no longer sets the poll interval, because that test finds the server on the first poll. Co-Authored-By: Claude Opus 5.5 --- internal/awsdevops/awsdevops_test.go | 11 ++++++-- internal/awsdevops/runplan_test.go | 42 ++++++++++++++++++++++++++-- 2 files changed, 47 insertions(+), 6 deletions(-) diff --git a/internal/awsdevops/awsdevops_test.go b/internal/awsdevops/awsdevops_test.go index b5220188e..80ef0e5c0 100644 --- a/internal/awsdevops/awsdevops_test.go +++ b/internal/awsdevops/awsdevops_test.go @@ -46,9 +46,11 @@ type fakeAgent struct { // operatorApps holds enabled apps by space. operatorApps map[string]*devopsagent.GetOperatorAppOutput // operatorAppErr fails every GetOperatorApp call. - operatorAppErr error - calls []string - services []agenttypes.RegisteredService + operatorAppErr error + calls []string + services []agenttypes.RegisteredService + // onListServices runs before ListServices returns. + onListServices func() associationInputs []*devopsagent.AssociateServiceInput registerOutput *devopsagent.RegisterServiceOutput registerErr error @@ -278,6 +280,9 @@ func (f *fakeAgent) ListAgentSpaces(_ context.Context, in *devopsagent.ListAgent func (f *fakeAgent) ListServices(_ context.Context, in *devopsagent.ListServicesInput, _ ...func(*devopsagent.Options)) (*devopsagent.ListServicesOutput, error) { f.calls = append(f.calls, "ListServices") + if f.onListServices != nil { + f.onListServices() + } items, next := page(f.services, in.NextToken, f.paginate) return &devopsagent.ListServicesOutput{Services: items, NextToken: next}, nil diff --git a/internal/awsdevops/runplan_test.go b/internal/awsdevops/runplan_test.go index 608b2f46c..cb1dc66a9 100644 --- a/internal/awsdevops/runplan_test.go +++ b/internal/awsdevops/runplan_test.go @@ -913,13 +913,16 @@ func oauthMCPService(serviceID string) agenttypes.RegisteredService { return service } -func TestOAuthRegistersAfterApproval(t *testing.T) { - defer awsdevops.SetOAuthPollInterval(time.Millisecond)() - agent := &fakeAgent{registerOutput: &devopsagent.RegisterServiceOutput{ +func oauthStepAgent() *fakeAgent { + return &fakeAgent{registerOutput: &devopsagent.RegisterServiceOutput{ AdditionalStep: &agenttypes.AdditionalServiceRegistrationStepMemberOauth{ Value: agenttypes.OAuthAdditionalStepDetails{AuthorizationUrl: aws.String("https://example.com/approve")}, }, }} +} + +func TestOAuthRegistersAfterApproval(t *testing.T) { + agent := oauthStepAgent() clients := newTestClients(agent, newFakeIAM()) opts := runPlanOptions() var opened string @@ -946,6 +949,39 @@ func TestOAuthRegistersAfterApproval(t *testing.T) { assert.Contains(t, agent.calls, "AssociateService") } +// The approval takes time, so setup polls until AWS registers the server. +func TestOAuthPollsUntilTheServerAppears(t *testing.T) { + defer awsdevops.SetOAuthPollInterval(time.Millisecond)() + agent := oauthStepAgent() + clients := newTestClients(agent, newFakeIAM()) + opts := runPlanOptions() + opened := false + opts.OpenURL = func(string) error { + opened = true + + return nil + } + polls := 0 + agent.onListServices = func() { + if !opened { + return + } + polls++ + if polls == 3 { + agent.services = append(agent.services, oauthMCPService("mcp-account")) + } + } + plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) + require.NoError(t, err) + + result, err := awsdevops.RunPlan(context.Background(), clients, opts, plan) + + require.NoError(t, err) + assert.Equal(t, 3, polls) + assert.Equal(t, "mcp-account", result.MCPServiceID) + assert.Contains(t, agent.calls, "AssociateService") +} + func TestOAuthReusesTheServerOfTheAccount(t *testing.T) { agent := &fakeAgent{ agentSpaces: []agenttypes.AgentSpace{namedSpace("space-7", "platform")}, From c6cecb86df7f3e026b96aa98332ac5ca73d49af0 Mon Sep 17 00:00:00 2001 From: Christina Chu Date: Wed, 28 Oct 2026 13:04:23 -0700 Subject: [PATCH 3/6] fix: stop when the AWS account has more than one OAuth MCP server Setup took the first OAuth server that ListServices sent. Each server acts as the member who approved it, so two agent spaces could act as different members. Now the plan stops, names each server, and asks for --mcp-service-id. Co-Authored-By: Claude Opus 5.5 --- internal/awsdevops/lookup.go | 22 +++++++++++++++++---- internal/awsdevops/runplan_test.go | 31 ++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 4 deletions(-) diff --git a/internal/awsdevops/lookup.go b/internal/awsdevops/lookup.go index d16c9434a..7ea6cf365 100644 --- a/internal/awsdevops/lookup.go +++ b/internal/awsdevops/lookup.go @@ -200,19 +200,33 @@ func agentSpaceARN(region, accountID, spaceID string) string { // accountMCPServer finds the MCP server of the AWS account: a server at // endpoint that signs in with OAuth. A server with a token is never a match, // even with the name MCPServerName, because its token belongs to one -// LaunchDarkly account and not to the AWS account. +// LaunchDarkly account and not to the AWS account. Each server acts as the +// member who approved it, so with more than one, the user must pick one. func accountMCPServer(ctx context.Context, clients Clients, endpoint string) (mcpRef, error) { services, err := listServices(ctx, clients) if err != nil { return mcpRef{}, err } + var ( + found mcpRef + all []string + ) for _, service := range services { - if mcpEndpointOf(service) == endpoint && usesOAuth(service) { - return mcpRef{ID: aws.ToString(service.ServiceId), Name: serviceName(service)}, nil + if mcpEndpointOf(service) != endpoint || !usesOAuth(service) { + continue } + found = mcpRef{ID: aws.ToString(service.ServiceId), Name: serviceName(service)} + all = append(all, fmt.Sprintf("%s (%s)", found.Name, found.ID)) + } + if len(all) > 1 { + return mcpRef{}, fmt.Errorf( + "the AWS account has %d LaunchDarkly MCP servers that use OAuth: %s. "+ + "To pick one, pass --mcp-service-id ", + len(all), strings.Join(all, ", "), + ) } - return mcpRef{}, nil + return found, nil } func usesOAuth(service agenttypes.RegisteredService) bool { diff --git a/internal/awsdevops/runplan_test.go b/internal/awsdevops/runplan_test.go index cb1dc66a9..b78807249 100644 --- a/internal/awsdevops/runplan_test.go +++ b/internal/awsdevops/runplan_test.go @@ -1015,3 +1015,34 @@ func TestOAuthDoesNotReuseAServerWithAToken(t *testing.T) { require.NoError(t, err) assert.Equal(t, awsdevops.MCPRegister, plan.MCP.Action) } + +// Each OAuth server acts as the member who approved it, so setup does not +// pick between two of them. +func TestPlanStopsOnAnAccountWithTwoOAuthServers(t *testing.T) { + agent := &fakeAgent{services: []agenttypes.RegisteredService{oauthMCPService("mcp-a"), oauthMCPService("mcp-b")}} + clients := newTestClients(agent, newFakeIAM()) + + _, err := awsdevops.PlanSetup(context.Background(), clients, runPlanOptions(), true) + + require.Error(t, err) + assert.ErrorContains(t, err, "the AWS account has 2 LaunchDarkly MCP servers that use OAuth") + assert.ErrorContains(t, err, "(mcp-a)") + assert.ErrorContains(t, err, "(mcp-b)") + assert.ErrorContains(t, err, "--mcp-service-id") +} + +func TestMCPServiceIDPicksOneOfTwoOAuthServers(t *testing.T) { + agent := &fakeAgent{services: []agenttypes.RegisteredService{oauthMCPService("mcp-a"), oauthMCPService("mcp-b")}} + clients := newTestClients(agent, newFakeIAM()) + opts := runPlanOptions() + opts.MCPServiceID = "mcp-b" + + plan, err := awsdevops.PlanSetup(context.Background(), clients, opts, true) + require.NoError(t, err) + require.Equal(t, awsdevops.MCPReuse, plan.MCP.Action) + result, err := awsdevops.RunPlan(context.Background(), clients, opts, plan) + + require.NoError(t, err) + assert.Equal(t, "mcp-b", result.MCPServiceID) + assert.NotContains(t, agent.calls, "RegisterService") +} From 71384b9f6fd0d6085be0ef3a82124477b8c4433b Mon Sep 17 00:00:00 2001 From: Christina Chu Date: Wed, 28 Oct 2026 13:08:48 -0700 Subject: [PATCH 4/6] fix: say what to do when the MCP server name is taken If RegisterService reports that LaunchDarkly-OAuth already exists, setup now says to rename or delete that server in the AWS console, or to pass --mcp-service-id. Before, setup showed only the error from AWS. Co-Authored-By: Claude Opus 5.5 --- internal/awsdevops/runplan_test.go | 19 +++++++++++++++++++ internal/awsdevops/setup.go | 10 ++++++++++ 2 files changed, 29 insertions(+) diff --git a/internal/awsdevops/runplan_test.go b/internal/awsdevops/runplan_test.go index b78807249..5d7f288be 100644 --- a/internal/awsdevops/runplan_test.go +++ b/internal/awsdevops/runplan_test.go @@ -1046,3 +1046,22 @@ func TestMCPServiceIDPicksOneOfTwoOAuthServers(t *testing.T) { assert.Equal(t, "mcp-b", result.MCPServiceID) assert.NotContains(t, agent.calls, "RegisterService") } + +// A server with the name and a token passes the plan, so the error from AWS +// must say what to do. +func TestOAuthNameClashSaysWhatToDo(t *testing.T) { + agent := &fakeAgent{registerErr: errors.New("ConflictException: a service named LaunchDarkly-OAuth already exists")} + clients := newTestClients(agent, newFakeIAM()) + plan, err := awsdevops.PlanSetup(context.Background(), clients, runPlanOptions(), true) + require.NoError(t, err) + + _, err = awsdevops.RunPlan(context.Background(), clients, runPlanOptions(), plan) + + require.Error(t, err) + assert.ErrorContains(t, err, "an MCP server named LaunchDarkly-OAuth is already registered") + assert.ErrorContains(t, err, "Rename or delete that server in the AWS console") + assert.ErrorContains(t, err, "--mcp-service-id ") + assert.ErrorContains(t, err, "ConflictException") + // Only the AWS account association, and no MCP server association. + assert.Equal(t, awsAssociationInputs(agent), agent.associationInputs) +} diff --git a/internal/awsdevops/setup.go b/internal/awsdevops/setup.go index cedc3bebb..d19fbd5c3 100644 --- a/internal/awsdevops/setup.go +++ b/internal/awsdevops/setup.go @@ -1202,6 +1202,16 @@ func registerOAuthMCPServer( }, } registration, err := registerWithRetry(ctx, clients, input, logf, endpoint) + if err != nil && strings.Contains(err.Error(), "already exists") { + // Setup reuses only an OAuth server at the endpoint, so this server + // uses a token or another endpoint, or it appeared after the check. + return fmt.Errorf( + "an MCP server named %s is already registered on this AWS account, so setup cannot register one. "+ + "Rename or delete that server in the AWS console, then run setup again. "+ + "Or, to connect a server that is already registered, pass --mcp-service-id : %w", + MCPOAuthServerName, err, + ) + } if err != nil { return fmt.Errorf("unable to register the LaunchDarkly MCP server: %w", err) } From 960d89b38530eb16915b49c6191c9d5571e6bb25 Mon Sep 17 00:00:00 2001 From: Christina Chu Date: Wed, 28 Oct 2026 13:12:15 -0700 Subject: [PATCH 5/6] docs: say that setup does not need the AWS console sign-in The help said that AWS can ask for a console sign-in before the approval. AWS registers the server before it opens the console, so the sign-in comes after the approval, and setup does not wait for it. Co-Authored-By: Claude Opus 5.5 --- cmd/aws_devops_agent/setup.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cmd/aws_devops_agent/setup.go b/cmd/aws_devops_agent/setup.go index 08085bf1e..3053258bd 100644 --- a/cmd/aws_devops_agent/setup.go +++ b/cmd/aws_devops_agent/setup.go @@ -51,7 +51,8 @@ operator app URL, where the agent is used. The AWS account gets one LaunchDarkly MCP server, which all agent spaces use. If the account has no server, setup registers one with OAuth. It opens a browser, where you approve access to LaunchDarkly, and waits until AWS -registers the server. AWS can ask you to sign in to the AWS console first. +registers the server. After the approval, the browser opens the AWS console, +which can ask you to sign in. Setup does not need that sign-in. The agent then acts as the LaunchDarkly member who approved access, in every agent space. If the agent space already has an MCP server, setup keeps it. From 56eca353af93236aef3e44fc7495c6021d08be9e Mon Sep 17 00:00:00 2001 From: Christina Chu Date: Wed, 28 Oct 2026 13:13:50 -0700 Subject: [PATCH 6/6] docs: remove the sentence about the console sign-in Co-Authored-By: Claude Opus 5.5 --- cmd/aws_devops_agent/setup.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cmd/aws_devops_agent/setup.go b/cmd/aws_devops_agent/setup.go index 3053258bd..0720bb69e 100644 --- a/cmd/aws_devops_agent/setup.go +++ b/cmd/aws_devops_agent/setup.go @@ -52,9 +52,9 @@ The AWS account gets one LaunchDarkly MCP server, which all agent spaces use. If the account has no server, setup registers one with OAuth. It opens a browser, where you approve access to LaunchDarkly, and waits until AWS registers the server. After the approval, the browser opens the AWS console, -which can ask you to sign in. Setup does not need that sign-in. -The agent then acts as the LaunchDarkly member who approved access, in every -agent space. If the agent space already has an MCP server, setup keeps it. +which can ask you to sign in. The agent then acts as the LaunchDarkly member +who approved access, in every agent space. If the agent space already has an +MCP server, setup keeps it. A space that reuses a server keeps its tool allowlist. If the allowlist misses tools that this ldcli version allows, setup warns. To add them, pass