override suspend fun connect()

in plugins/toolkit/jetbrains-gateway/src/software/aws/toolkits/jetbrains/gateway/CawsConnectionProvider.kt [102:388]


    override suspend fun connect(parameters: Map<String, String>, requestor: ConnectionRequestor): GatewayConnectionHandle? {
        val connectionParams = try {
            CawsConnectionParameters.fromParameters(parameters)
        } catch (e: Exception) {
            LOG.error(e) { "Caught exception while building connection settings" }
            Messages.showErrorDialog(e.message ?: message("general.unknown_error"), message("caws.workspace.connection.failed"))
            return null
        }

        val currentConnection = service<ToolkitConnectionManager>().activeConnectionForFeature(CodeCatalystConnection.getInstance())
            as AwsBearerTokenConnection?

        val ssoSettings = connectionParams.ssoSettings ?: SsoSettings(SONO_URL, SONO_REGION)

        if (currentConnection != null) {
            if (ssoSettings.startUrl != currentConnection.startUrl) {
                val ans = Messages.showOkCancelDialog(
                    message("gateway.auth.different.account.required", ssoSettings.startUrl),
                    message("gateway.auth.different.account.sign.in"),
                    message("caws.login"),
                    message("general.cancel"),
                    Messages.getErrorIcon(),
                    null
                )
                if (ans == Messages.OK) {
                    logoutFromSsoConnection(project = null, currentConnection)
                    loginSso(project = null, ssoSettings.startUrl, ssoSettings.region, CODECATALYST_SCOPES)
                } else {
                    return null
                }
            }
        }

        val connectionSettings = try {
            CodeCatalystCredentialManager.getInstance().getConnectionSettings() ?: error("Unable to find connection settings")
        } catch (e: ProcessCanceledException) {
            return null
        }

        val userId = lazilyGetUserId()

        val spaceName = connectionParams.space
        val projectName = connectionParams.project
        val envId = connectionParams.envId
        val id = WorkspaceIdentifier(CawsProject(spaceName, projectName), envId)

        val lifetime = Lifetime.Eternal.createNested()
        val workflowDisposable = Lifetime.Eternal.createNestedDisposable()

        return CawsGatewayConnectionHandle(lifetime, envId) {
            // reference lost with all the blocks
            it.let { gatewayHandle ->
                val view = JBTabbedPane()
                val workflowEmitter = TabbedWorkflowEmitter(view, workflowDisposable)

                fun handleException(e: Throwable) {
                    if (e is ProcessCanceledException || e is CancellationException) {
                        CodecatalystTelemetry.connect(project = null, userId = userId, result = TelemetryResult.Cancelled)
                        LOG.warn { "Connect to dev environment cancelled" }
                    } else {
                        CodecatalystTelemetry.connect(project = null, userId = userId, result = TelemetryResult.Failed, reason = e.javaClass.simpleName)
                        LOG.error(e) { "Caught exception while connecting to dev environment" }
                    }
                    lifetime.terminate()
                }

                // TODO: Describe env to validate JB ide is set on it
                lifetime.launch {
                    try {
                        val cawsClient = connectionSettings.awsClient<CodeCatalystClient>()
                        val environmentActions = WorkspaceActions(spaceName, projectName, envId, cawsClient)
                        val executor = CawsCommandExecutor(cawsClient, envId, spaceName, projectName)

                        // should probably consider logging output to logger as well
                        // on failure we should display meaningful error and put retry button somewhere
                        lifetime.startUnderModalProgressAsync(
                            title = message("caws.connecting.waiting_for_environment"),
                            canBeCancelled = true,
                            isIndeterminate = true,
                        ) {
                            val timeBeforeEnvIsRunningCheck = System.currentTimeMillis()
                            var validateEnvIsRunningResult = TelemetryResult.Succeeded
                            var errorMessageDuringStateValidation: String? = null
                            try {
                                validateEnvironmentIsRunning(indicator, environmentActions)
                            } catch (e: Exception) {
                                validateEnvIsRunningResult = TelemetryResult.Failed
                                errorMessageDuringStateValidation = e.message
                                throw e
                            } finally {
                                CodecatalystTelemetry.devEnvironmentWorkflowStatistic(
                                    project = null,
                                    userId = userId,
                                    result = validateEnvIsRunningResult,
                                    duration = (System.currentTimeMillis() - timeBeforeEnvIsRunningCheck).toDouble(),
                                    codecatalystDevEnvironmentWorkflowStep = "validateEnvRunning",
                                    codecatalystDevEnvironmentWorkflowError = errorMessageDuringStateValidation
                                )
                            }

                            lifetime.launchIOBackground {
                                ApplicationManager.getApplication().messageBus.syncPublisher(WorkspaceNotifications.TOPIC)
                                    .environmentStarted(
                                        WorkspaceListStateChangeContext(
                                            WorkspaceIdentifier(CawsProject(spaceName, projectName), envId)
                                        )
                                    )
                            }

                            val pluginPath = "$IDE_BACKEND_DIR/plugins/${AwsToolkit.PLUGINS_INFO.getValue(AwsPlugin.TOOLKIT).path?.fileName}"
                            var retries = 3
                            val startTimeToCheckInstallation = System.currentTimeMillis()

                            val toolkitInstallSettings: ToolkitInstallSettings? = coroutineScope {
                                while (retries > 0) {
                                    indicator.checkCanceled()
                                    val pluginIsInstalled = executor.remoteDirectoryExists(
                                        pluginPath,
                                        timeout = Duration.ofSeconds(15)
                                    )

                                    when (pluginIsInstalled) {
                                        null -> {
                                            if (retries == 1) {
                                                return@coroutineScope null
                                            } else {
                                                retries--
                                                continue
                                            }
                                        }

                                        true -> return@coroutineScope ToolkitInstallSettings.None
                                        false -> return@coroutineScope connectionParams.toolkitInstallSettings
                                    }
                                }
                            } as ToolkitInstallSettings?

                            toolkitInstallSettings ?: let {
                                // environment is non-responsive to SSM; restart
                                LOG.warn { "Restarting $envId since it appears unresponsive to SSM Run-Command" }
                                val timeTakenToCheckInstallation = System.currentTimeMillis() - startTimeToCheckInstallation
                                CodecatalystTelemetry.devEnvironmentWorkflowStatistic(
                                    project = null,
                                    userId = userId,
                                    result = TelemetryResult.Failed,
                                    codecatalystDevEnvironmentWorkflowStep = "ToolkitInstallationSSMCheck",
                                    codecatalystDevEnvironmentWorkflowError = "Timeout/Unknown error while connecting to Dev Env via SSM",
                                    duration = timeTakenToCheckInstallation.toDouble()
                                )

                                launchChildSyncIOBackground {
                                    environmentActions.stopEnvironment()
                                    GatewayUI.getInstance().connect(parameters)
                                }

                                gatewayHandle.terminate()
                                return@startUnderModalProgressAsync JLabel()
                            }

                            lifetime.startUnderBackgroundProgressAsync(message("caws.download.thin_client"), isIndeterminate = true) {
                                val (backendVersion, getBackendVersionTime) = measureTimedValue {
                                    tryOrNull {
                                        executor.executeCommandNonInteractive(
                                            "sh",
                                            "-c",
                                            GET_IDE_BACKEND_VERSION_COMMAND,
                                            timeout = Duration.ofSeconds(15)
                                        ).stdout
                                    }
                                }
                                CodecatalystTelemetry.devEnvironmentWorkflowStatistic(
                                    project = null,
                                    userId = userId,
                                    result = if (backendVersion != null) TelemetryResult.Succeeded else TelemetryResult.Failed,
                                    duration = getBackendVersionTime.toDouble(DurationUnit.MILLISECONDS),
                                    codecatalystDevEnvironmentWorkflowStep = "getBackendVersion"
                                )

                                if (backendVersion.isNullOrBlank()) {
                                    LOG.warn { "Could not determine backend version to prefetch thin client" }
                                } else {
                                    val (clientPaths, downloadClientTime) = measureTimedValue {
                                        BuildNumber.fromStringOrNull(backendVersion)?.asStringWithoutProductCode()?.let { build ->
                                            LOG.info { "Fetching client for version: $build" }
                                            CodeWithMeClientDownloader.downloadClientAndJdk(build, indicator)
                                        }
                                    }

                                    CodecatalystTelemetry.devEnvironmentWorkflowStatistic(
                                        project = null,
                                        userId = userId,
                                        result = if (clientPaths != null) TelemetryResult.Succeeded else TelemetryResult.Failed,
                                        duration = downloadClientTime.toDouble(DurationUnit.MILLISECONDS),
                                        codecatalystDevEnvironmentWorkflowStep = "downloadThinClient"
                                    )
                                }
                            }

                            runBackendWorkflow(
                                view,
                                workflowEmitter,
                                userId,
                                indicator,
                                lifetime.createNested(),
                                parameters,
                                executor,
                                id,
                                connectionParams.gitSettings,
                                toolkitInstallSettings
                            ).await()
                        }.invokeOnCompletion { e ->
                            if (e == null) {
                                CodecatalystTelemetry.connect(project = null, userId = userId, result = TelemetryResult.Succeeded)
                                lifetime.onTermination {
                                    Disposer.dispose(workflowDisposable)
                                }
                            } else {
                                handleException(e)
                                if (e is ProcessCanceledException || e is CancellationException) {
                                    return@invokeOnCompletion
                                }
                                runInEdt {
                                    DialogBuilder().apply {
                                        setCenterPanel(
                                            panel {
                                                row {
                                                    icon(AllIcons.General.ErrorDialog).align(AlignY.TOP)

                                                    panel {
                                                        row {
                                                            label(message("caws.workspace.connection.failed")).applyToComponent {
                                                                font = JBFont.regular().asBold()
                                                            }
                                                        }

                                                        row {
                                                            label(e.message ?: message("general.unknown_error"))
                                                        }
                                                    }
                                                }

                                                if (view.tabCount != 0) {
                                                    collapsibleGroup(message("general.logs"), false) {
                                                        row {
                                                            cell(view)
                                                                .align(AlignX.FILL)
                                                        }
                                                    }.expanded = false
                                                    // TODO: can't seem to reliably force a terminal redraw on initial expand
                                                }
                                            }
                                        )

                                        addOkAction()
                                        addCancelAction()
                                        okAction.setText(message("settings.retry"))
                                        setOkOperation {
                                            dialogWrapper.close(DialogWrapper.OK_EXIT_CODE)
                                            GatewayUI.getInstance().connect(parameters)
                                        }
                                    }.show()
                                    Disposer.dispose(workflowDisposable)
                                }
                            }
                        }
                    } catch (e: Exception) {
                        handleException(e)
                        if (e is ProcessCanceledException || e is CancellationException) {
                            return@launch
                        }

                        runInEdt {
                            Messages.showErrorDialog(e.message ?: message("general.unknown_error"), message("caws.workspace.connection.failed"))
                        }
                        throw e
                    }
                }

                return@let panel {
                    row {
                        cell(view)
                            .align(Align.FILL)
                    }
                }
            }
        }
    }