Merge remote-tracking branch 'upstream/main' into multi-window

Author: Will-ShaoHua

.changes/next-release/bugfix-a70cb565-5de8-4302-9f1e-71925f78061b.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Amazon Q: Customization now resets with a warning if unavailable in the selected profile."
+}

.changes/next-release/bugfix-f31946b6-3709-44a5-b23f-9fa202bd5cb5.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Fix infinite loop when workspace indexing server fails to initialize"
+}

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt

@@ -22,6 +22,8 @@ import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
 import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
+import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
+import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
 import software.aws.toolkits.jetbrains.utils.notifyInfo
@@ -48,6 +50,7 @@ private fun notifyInvalidSelectedCustomization(project: Project) {
 }
 
 private fun notifyNewCustomization(project: Project) {
+    if (ApplicationManager.getApplication().isUnitTestMode) return
     notifyInfo(
         title = message("codewhisperer.custom.dialog.title"),
         content = message("codewhisperer.notification.custom.new_customization"),
@@ -81,6 +84,18 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
 
     private var customizationArnOverrideV2: String? = null
 
+    init {
+        ApplicationManager.getApplication().messageBus.connect(this).subscribe(
+            QRegionProfileSelectedListener.TOPIC,
+            object : QRegionProfileSelectedListener {
+                override fun onProfileSelected(project: Project, profile: QRegionProfile?) {
+                    pluginAwareExecuteOnPooledThread {
+                        CodeWhispererModelConfigurator.getInstance().listCustomizations(project, passive = true)
+                    }
+                }
+            }
+        )
+    }
     override fun showConfigDialog(project: Project) {
         runInEdt {
             calculateIfIamIdentityCenterConnection(project) {

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt

@@ -7,8 +7,10 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.testFramework.ApplicationRule
 import com.intellij.testFramework.DisposableRule
 import com.intellij.testFramework.ProjectRule
+import com.intellij.testFramework.registerServiceInstance
 import com.intellij.testFramework.replaceService
 import com.intellij.util.xmlb.XmlSerializer
+import migration.software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import org.assertj.core.api.Assertions.assertThat
 import org.jdom.output.XMLOutputter
 import org.junit.Before
@@ -40,10 +42,14 @@ import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
 import software.aws.toolkits.jetbrains.core.region.MockRegionProviderRule
 import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
 import software.aws.toolkits.jetbrains.services.amazonq.FeatureContext
+import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
+import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomizationState
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.DefaultCodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.utils.xmlElement
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
 import kotlin.reflect.full.memberProperties
 import kotlin.reflect.jvm.isAccessible
 
@@ -75,6 +81,7 @@ class CodeWhispererModelConfiguratorTest {
     private lateinit var sut: DefaultCodeWhispererModelConfigurator
     private lateinit var mockClient: CodeWhispererRuntimeClient
     private lateinit var abManager: CodeWhispererFeatureConfigService
+    private lateinit var mockClintAdaptor: CodeWhispererClientAdaptor
 
     @Before
     fun setup() {
@@ -83,7 +90,11 @@ class CodeWhispererModelConfiguratorTest {
         regionProvider.addRegion(Region.US_EAST_1)
         regionProvider.addRegion(Region.US_EAST_2)
 
-        sut = DefaultCodeWhispererModelConfigurator()
+        sut = spy(CodeWhispererModelConfigurator.getInstance() as DefaultCodeWhispererModelConfigurator).also { spyInstance ->
+            ApplicationManager.getApplication().replaceService(
+                DefaultCodeWhispererModelConfigurator::class.java, spyInstance, disposableRule.disposable
+            )
+        }
 
         (ToolkitConnectionManager.getInstance(projectRule.project) as DefaultToolkitConnectionManager).loadState(ToolkitConnectionManagerState())
         mockClient.stub {
@@ -110,6 +121,9 @@ class CodeWhispererModelConfiguratorTest {
             abManager,
             disposableRule.disposable
         )
+
+        mockClintAdaptor = mock()
+        projectRule.project.registerServiceInstance(CodeWhispererClientAdaptor::class.java, mockClintAdaptor)
     }
 
     @Test
@@ -550,4 +564,48 @@ class CodeWhispererModelConfiguratorTest {
         assertThat(actual.previousAvailableCustomizations).hasSize(1)
         assertThat(actual.previousAvailableCustomizations["fake-sso-url"]).isEqualTo(listOf("arn_1", "arn_2", "arn_3"))
     }
+
+    @Test
+    fun `profile switch should keep using existing customization if new list still contains that arn`() {
+        val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url 1", scopes = Q_SCOPES))
+        ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn)
+        val oldCustomization = CodeWhispererCustomization("oldArn", "oldName", "oldDescription")
+        sut.switchCustomization(projectRule.project, oldCustomization)
+
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
+
+        val fakeCustomizations = listOf(
+            CodeWhispererCustomization("oldArn", "oldName", "oldDescription")
+        )
+        mockClintAdaptor.stub { on { listAvailableCustomizations() } doReturn fakeCustomizations }
+
+        ApplicationManager.getApplication().messageBus
+            .syncPublisher(QRegionProfileSelectedListener.TOPIC)
+            .onProfileSelected(projectRule.project, null)
+
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
+    }
+
+    @Test
+    fun `profile switch should invalidate obsolete customization if it's not in the new list`() {
+        val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url 1", scopes = Q_SCOPES))
+        ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn)
+        val oldCustomization = CodeWhispererCustomization("oldArn", "oldName", "oldDescription")
+        sut.switchCustomization(projectRule.project, oldCustomization)
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
+        val fakeCustomizations = listOf(
+            CodeWhispererCustomization("newArn", "newName", "newDescription")
+        )
+        mockClintAdaptor.stub { on { listAvailableCustomizations() } doReturn fakeCustomizations }
+
+        val latch = CountDownLatch(1)
+
+        ApplicationManager.getApplication().messageBus
+            .syncPublisher(QRegionProfileSelectedListener.TOPIC)
+            .onProfileSelected(projectRule.project, null)
+
+        latch.await(2, TimeUnit.SECONDS)
+
+        assertThat(sut.activeCustomization(projectRule.project)).isNull()
+    }
 }

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt

@@ -116,10 +116,14 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
                     logger.info { "project context index starting" }
                     delay(300)
                     val isIndexSuccess = index()
-                    if (isIndexSuccess) isIndexComplete.set(true)
+                    if (isIndexSuccess) {
+                        isIndexComplete.set(true)
+                    }
                     return
                 }
+                retryCount.incrementAndGet()
             } catch (e: Exception) {
+                logger.warn(e) { "failed to init project context" }
                 if (e.stackTraceToString().contains("Connection refused")) {
                     retryCount.incrementAndGet()
                     delay(10000)
@@ -133,6 +137,7 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
     private suspend fun initEncryption(): Boolean {
         val request = encoderServer.getEncryptionRequest()
         val response = sendMsgToLsp(LspMessage.Initialize, request)
+        logger.info { "received response to init encryption: $response" }
         return response?.responseCode == 200
     }
 
@@ -315,12 +320,12 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
     }
 
     private suspend fun sendMsgToLsp(msgType: LspMessage, request: String?): LspResponse? {
-        logger.info { "sending message: ${msgType.endpoint} to lsp on port ${encoderServer.port}" }
-        val url = URI("http://127.0.0.1:${encoderServer.port}/${msgType.endpoint}").toURL()
         if (!encoderServer.isNodeProcessRunning()) {
             logger.warn { "language server for ${project.name} is not running" }
             return null
         }
+        logger.info { "sending message: ${msgType.endpoint} to lsp on port ${encoderServer.port}" }
+        val url = URI("http://127.0.0.1:${encoderServer.port}/${msgType.endpoint}").toURL()
         // use 1h as timeout for index, 5 seconds for other APIs
         val timeoutMs = if (msgType is LspMessage.Index) 60.minutes.inWholeMilliseconds.toInt() else 5000
         // dedicate single thread to index operation because it can be long running