Update authentication logic to handle network issues

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAuthManager.kt

@@ -3,6 +3,7 @@
 
 package software.aws.toolkits.jetbrains.core.credentials
 
+import com.intellij.ide.util.PropertiesComponent
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.components.service
 import com.intellij.openapi.extensions.ExtensionPointName
@@ -25,12 +26,15 @@
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.InteractiveBearerTokenProvider
+import software.aws.toolkits.jetbrains.utils.notifyInfo
 import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
 import software.aws.toolkits.resources.AwsCoreBundle
+import software.aws.toolkits.resources.AwsCoreBundle.message
 import software.aws.toolkits.telemetry.AuthTelemetry
 import software.aws.toolkits.telemetry.CredentialSourceId
 import software.aws.toolkits.telemetry.CredentialType
 import software.aws.toolkits.telemetry.Result
+import java.net.UnknownHostException
 import java.time.Instant
 
 sealed interface ToolkitConnection {
@@ -254,7 +258,7 @@
     var didReauth = false
     maybeReauthProviderIfNeeded(project, reauthSource, tokenProvider) {
         didReauth = true
         runUnderProgressIfNeeded(project, AwsCoreBundle.message("credentials.pending.title"), true) {
             try {
                 tokenProvider.reauthenticate()
                 if (isReAuth) {
@@ -272,6 +276,7 @@
                         source = source,
                     )
                 }
+                hasSeenFirstNetworkError = false
             } catch (e: Exception) {
                 if (isReAuth) {
                     val result = if (e is ProcessCanceledException) Result.Cancelled else Result.Failed
@@ -320,16 +325,36 @@
         BearerTokenAuthState.NEEDS_REFRESH -> {
             try {
                 getLogger<ToolkitAuthManager>().warn { "Starting token refresh" }
                 return runUnderProgressIfNeeded(project, AwsCoreBundle.message("credentials.refreshing"), true) {
                     tokenProvider.resolveToken()
                     BearerTokenProviderListener.notifyCredUpdate(tokenProvider.id)
+                    hasSeenFirstNetworkError = false
                     return@runUnderProgressIfNeeded false
                 }
-            } catch (e: SsoOidcException) {
-                AuthTelemetry.sourceOfRefresh(authRefreshSource = reauthSource.toString())
-                getLogger<ToolkitAuthManager>().warn(e) { "Redriving bearer token login flow since token could not be refreshed" }
-                onReauthRequired(e)
-                return true
+            } catch (e: Exception) {
+                when {
+                    e is SsoOidcException -> {
+                        AuthTelemetry.sourceOfRefresh(authRefreshSource = reauthSource.toString())
+                        getLogger<ToolkitAuthManager>().warn(e) { "Redriving bearer token login flow since token could not be refreshed" }
+                        onReauthRequired(e)
+                        return true
+                    }
+                    e is UnknownHostException || e.message?.contains("Unable to execute HTTP request") == true -> {
+                        getLogger<ToolkitAuthManager>().warn(e) { "Failed to refresh token" }
+                        synchronized(networkErrorLock) {
+                            if (!hasSeenFirstNetworkError) {
+                                hasSeenFirstNetworkError = true
+                                notifyInfo(
+                                    message("general.auth.network.error"),
+                                    message("general.auth.network.error.message"),
+                                    project
+                                )
+                            }
+                        }
+                        return false
+                    }
+                    else -> { return false }
+                }
             }
         }
 
@@ -410,6 +435,9 @@
     }
 }
 
+private var hasSeenFirstNetworkError = false
+private val networkErrorLock = Object()
+
 data class ConnectionMetadata(
     val sourceId: String? = null,
 )

plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/ToolkitAuthManagerTest.kt

@@ -0,0 +1,169 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.core.credentials
+
+import com.intellij.ide.util.PropertiesComponent
+import com.intellij.openapi.project.Project
+import com.intellij.testFramework.ApplicationExtension
+import io.mockk.every
+import io.mockk.just
+import io.mockk.mockkObject
+import io.mockk.mockkStatic
+import io.mockk.runs
+import io.mockk.verify
+import org.junit.jupiter.api.Assertions.*
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.TestInstance
+import org.junit.jupiter.api.extension.ExtendWith
+import org.mockito.kotlin.doThrow
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.reset
+import org.mockito.kotlin.whenever
+import software.aws.toolkits.jetbrains.core.credentials.sso.DeviceAuthorizationGrantToken
+import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenAuthState
+import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
+import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
+import software.aws.toolkits.jetbrains.utils.notifyInfo
+import software.aws.toolkits.resources.AwsCoreBundle.message
+import java.time.Instant
+import java.time.temporal.ChronoUnit
+
+@ExtendWith(ApplicationExtension::class)
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+class ToolkitAuthManagerTest {
+    private lateinit var project: Project
+    private lateinit var tokenProvider: BearerTokenProvider
+    private var reauthCallCount = 0
+
+    @BeforeEach
+    fun setUp() {
+        project = mock()
+        tokenProvider = mock()
+        reauthCallCount = 0
+
+        val field = Class.forName("software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManagerKt")
+            .getDeclaredField("hasSeenFirstNetworkError")
+        field.isAccessible = true
+        field.set(null, false)
+
+        mockkObject(BearerTokenProviderListener)
+        mockkStatic("software.aws.toolkits.jetbrains.utils.NotificationUtilsKt")
+        every {
+            notifyInfo(any(), any(), any())
+        } just runs
+        every { BearerTokenProviderListener.notifyCredUpdate(any<String>()) } just runs
+    }
+
+    @Test
+    fun `test NEEDS_REFRESH state with network error - first occurrence`() {
+        whenever(tokenProvider.state()).thenReturn(BearerTokenAuthState.NEEDS_REFRESH)
+        doThrow(RuntimeException("Unable to execute HTTP request"))
+            .`when`(tokenProvider)
+            .resolveToken()
+
+        val result = maybeReauthProviderIfNeeded(
+            project,
+            ReauthSource.TOOLKIT,
+            tokenProvider
+        ) { _ -> reauthCallCount++ }
+
+        assertFalse(result)
+        assertEquals(0, reauthCallCount)
+        verify(exactly = 1) {
+            notifyInfo(
+                message("general.auth.network.error"),
+                message("general.auth.network.error.message"),
+                project
+            )
+        }
+    }
+
+    @Test
+    fun `test NEEDS_REFRESH state with network error - subsequent occurrence`() {
+        whenever(tokenProvider.state()).thenReturn(BearerTokenAuthState.NEEDS_REFRESH)
+        doThrow(RuntimeException("Unable to execute HTTP request"))
+            .`when`(tokenProvider)
+            .resolveToken()
+
+        // First call to set the internal flag
+        maybeReauthProviderIfNeeded(
+            project,
+            ReauthSource.TOOLKIT,
+            tokenProvider
+        ) { _ -> reauthCallCount++ }
+
+        // Second call - should not show notification
+        val result = maybeReauthProviderIfNeeded(
+            project,
+            ReauthSource.TOOLKIT,
+            tokenProvider
+        ) { _ -> reauthCallCount++ }
+
+        assertFalse(result)
+        assertEquals(0, reauthCallCount)
+        verify(exactly = 1) {
+            notifyInfo(
+                message("general.auth.network.error"),
+                message("general.auth.network.error.message"),
+                project
+            )
+        }
+    }
+
+    @Test
+    fun `test successful refresh clears notification flag`() {
+        whenever(tokenProvider.state()).thenReturn(BearerTokenAuthState.NEEDS_REFRESH)
+
+        // First trigger a network error
+        doThrow(RuntimeException("Unable to execute HTTP request"))
+            .`when`(tokenProvider)
+            .resolveToken()
+
+        maybeReauthProviderIfNeeded(
+            project,
+            ReauthSource.TOOLKIT,
+            tokenProvider
+        ) { _ -> reauthCallCount++ }
+
+        reset(tokenProvider)
+
+        // Now simulate successful refresh
+        whenever(tokenProvider.state()).thenReturn(BearerTokenAuthState.NEEDS_REFRESH)
+        whenever(tokenProvider.resolveToken()).thenReturn(
+            DeviceAuthorizationGrantToken(
+                startUrl = "https://example.com",
+                region = "us-east-1",
+                accessToken = "testAccessToken",
+                refreshToken = "testRefreshToken",
+                expiresAt = Instant.now().plus(1, ChronoUnit.HOURS),
+            )
+        )
+        maybeReauthProviderIfNeeded(
+            project,
+            ReauthSource.TOOLKIT,
+            tokenProvider
+        ) { _ -> reauthCallCount++ }
+
+        reset(tokenProvider)
+        // Now trigger another network error - should show notification again
+        whenever(tokenProvider.state()).thenReturn(BearerTokenAuthState.NEEDS_REFRESH)
+        doThrow(RuntimeException("Unable to execute HTTP request"))
+            .`when`(tokenProvider)
+            .resolveToken()
+        maybeReauthProviderIfNeeded(
+            project,
+            ReauthSource.TOOLKIT,
+            tokenProvider
+        ) { _ -> reauthCallCount++ }
+
+        verify(exactly = 2) {
+            notifyInfo(
+                message("general.auth.network.error"),
+                message("general.auth.network.error.message"),
+                project
+            )
+        }
+    }
+}

plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties

@@ -1240,6 +1240,8 @@ gateway.connection.workflow.step_skipped=Step skipped
 gateway.connection.workflow.step_successful=\nStep completed successfully\n
 general.add.another=Add another
 general.auth.reauthenticate=Reauthenticate
+general.auth.network.error=Network Error
+general.auth.network.error.message=Failed to refresh connection due to network issues
 general.cancel=Cancel
 general.close_button=Close
 general.configure_button=Configure