Add metric for location of reauth (#5021)

* Add metric for location of reauth * fixed tests * Removed comment
aws · Oct 29, 2024 · f8d7360 · f8d7360
1 parent 4eb2943
commit f8d7360
Show file tree

Hide file tree

Showing 16 changed files with 59 additions and 15 deletions.
diff --git a/...y/src/software/aws/toolkits/jetbrains/services/amazonq/explorerActions/SignInToQAction.kt b/...y/src/software/aws/toolkits/jetbrains/services/amazonq/explorerActions/SignInToQAction.kt
@@ -8,6 +8,7 @@ import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.project.DumbAwareAction
 import com.intellij.openapi.wm.ToolWindowManager
+import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
 import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
@@ -39,7 +40,7 @@ abstract class SignInToQActionBase(actionName: String) : DumbAwareAction(actionN
         UiTelemetry.click(project, "auth_start_Q")
         val connectionManager = ToolkitConnectionManager.getInstance(project)
         connectionManager.activeConnectionForFeature(QConnection.getInstance())?.let {
-            reauthConnectionIfNeeded(project, it, isReAuth = true)
+            reauthConnectionIfNeeded(project, it, isReAuth = true, reauthSource = ReauthSource.CODEWHISPERER_STATUSBAR)
         } ?: run {
             runInEdt {
                 if (requestCredentialsForQ(project, isReauth = false)) {

diff --git a/...nity/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt b/...nity/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt
@@ -26,6 +26,7 @@ import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
+import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.maybeReauthProviderIfNeeded
@@ -171,7 +172,7 @@ object CodeWhispererUtil {
         if (!isQExpired(project)) return false
         val tokenProvider = tokenProvider(project) ?: return false
         return try {
-            maybeReauthProviderIfNeeded(project, tokenProvider) {
+            maybeReauthProviderIfNeeded(project, ReauthSource.CODEWHISPERER, tokenProvider) {
                 runInEdt {
                     if (!CodeWhispererService.hasReAuthPromptBeenShown()) {
                         notifyConnectionExpiredRequestReauth(project)
@@ -256,7 +257,7 @@ object CodeWhispererUtil {
         val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())
         if (connection !is ManagedBearerSsoConnection) return
         pluginAwareExecuteOnPooledThread {
-            reauthConnectionIfNeeded(project, connection, isReAuth = true)
+            reauthConnectionIfNeeded(project, connection, isReAuth = true, reauthSource = ReauthSource.CODEWHISPERER)
         }
     }
 

diff --git a/.../tst/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtilTest.kt b/.../tst/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtilTest.kt
@@ -16,6 +16,7 @@ import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.extension.RegisterExtension
 import software.aws.toolkits.core.utils.test.aString
 import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
+import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
@@ -51,7 +52,7 @@ class CodeWhispererUtilTest {
         CodeWhispererUtil.reconnectCodeWhisperer(projectExtension.project)
 
         verify {
-            reauthConnectionIfNeeded(projectExtension.project, mockConnection, isReAuth = true)
+            reauthConnectionIfNeeded(projectExtension.project, mockConnection, isReAuth = true, reauthSource = ReauthSource.CODEWHISPERER)
         }
     }
 }
diff --git a/plugins/core/jetbrains-community/resources/telemetryOverride.json b/plugins/core/jetbrains-community/resources/telemetryOverride.json
@@ -273,6 +273,11 @@
             "type": "string",
             "description": "Comma-delimited list of enabled auths."
         },
+        {
+            "name": "authRefreshSource",
+            "type": "string",
+            "description": "Source triggering token refresh"
+        },
         {
             "name": "component",
             "allowedValues": [
@@ -813,6 +818,17 @@
                     "required": true
                 }
             ]
+        },
+        {
+            "name": "auth_sourceOfRefresh",
+            "description": "Source of user triggered refresh",
+            "metadata": [
+                {
+                    "type": "authRefreshSource",
+                    "required": false
+                }
+            ]
         }
+
     ]
 }
diff --git a/...re/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/LoginUtils.kt b/...re/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/LoginUtils.kt
@@ -198,7 +198,7 @@ fun authAndUpdateConfig(
 
     val connection = try {
         ToolkitAuthManager.getInstance().tryCreateTransientSsoConnection(updatedProfile) { connection ->
-            reauthConnectionIfNeeded(project, connection, onPendingToken)
+            reauthConnectionIfNeeded(project, connection, onPendingToken, reauthSource = ReauthSource.FRESH_AUTH)
         }
     } catch (e: Exception) {
         onError(e)

diff --git a/...ains-community/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAuthManager.kt b/...ains-community/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAuthManager.kt
@@ -27,6 +27,7 @@ import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenPr
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.InteractiveBearerTokenProvider
 import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
 import software.aws.toolkits.resources.AwsCoreBundle
+import software.aws.toolkits.telemetry.AuthTelemetry
 import software.aws.toolkits.telemetry.CredentialSourceId
 import software.aws.toolkits.telemetry.CredentialType
 import software.aws.toolkits.telemetry.Result
@@ -135,6 +136,7 @@ fun loginSso(
                     onPendingToken = onPendingToken,
                     isReAuth = false,
                     source = metadata?.sourceId,
+                    reauthSource = ReauthSource.FRESH_AUTH
                 )
             }
         } catch (e: Exception) {
@@ -182,6 +184,7 @@ fun loginSso(
             onPendingToken = onPendingToken,
             isReAuth = true,
             source = metadata?.sourceId,
+            reauthSource = ReauthSource.COMMON_LOGIN
         )
         return@let connection
     }
@@ -240,6 +243,7 @@ fun reauthConnectionIfNeeded(
     onPendingToken: (InteractiveBearerTokenProvider) -> Unit = {},
     isReAuth: Boolean = false,
     source: String? = null,
+    reauthSource: ReauthSource? = ReauthSource.TOOLKIT,
 ): BearerTokenProvider {
     val tokenProvider = (connection.getConnectionSettings() as TokenConnectionSettings).tokenProvider.delegate as BearerTokenProvider
     if (tokenProvider is InteractiveBearerTokenProvider) {
@@ -248,7 +252,7 @@ fun reauthConnectionIfNeeded(
 
     val startUrl = (connection as AwsBearerTokenConnection).startUrl
     var didReauth = false
-    maybeReauthProviderIfNeeded(project, tokenProvider) {
+    maybeReauthProviderIfNeeded(project, reauthSource, tokenProvider) {
         didReauth = true
         runUnderProgressIfNeeded(project, AwsCoreBundle.message("credentials.pending.title"), true) {
             try {
@@ -301,6 +305,7 @@ fun reauthConnectionIfNeeded(
 // Return true if need to re-auth, false otherwise
 fun maybeReauthProviderIfNeeded(
     project: Project?,
+    reauthSource: ReauthSource? = ReauthSource.TOOLKIT,
     tokenProvider: BearerTokenProvider,
     onReauthRequired: (SsoOidcException?) -> Any,
 ): Boolean {
@@ -314,12 +319,14 @@ fun maybeReauthProviderIfNeeded(
 
         BearerTokenAuthState.NEEDS_REFRESH -> {
             try {
+                getLogger<ToolkitAuthManager>().warn { "Starting token refresh" }
                 return runUnderProgressIfNeeded(project, AwsCoreBundle.message("credentials.refreshing"), true) {
                     tokenProvider.resolveToken()
                     BearerTokenProviderListener.notifyCredUpdate(tokenProvider.id)
                     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
@@ -332,6 +339,17 @@ fun maybeReauthProviderIfNeeded(
     }
 }
 
+enum class ReauthSource {
+    CODEWHISPERER,
+    TOOLKIT,
+    Q_CHAT,
+    LOGIN_BROWSER,
+    CODEWHISPERER_STATUSBAR,
+    CODECATALYST,
+    COMMON_LOGIN,
+    FRESH_AUTH,
+}
+
 fun deleteSsoConnection(connection: ProfileSsoManagedBearerSsoConnection) =
     deleteSsoConnection(connection.configSessionName)
 

diff --git a/...ware/aws/toolkits/jetbrains/core/credentials/profiles/ProfileCredentialProviderFactory.kt b/...ware/aws/toolkits/jetbrains/core/credentials/profiles/ProfileCredentialProviderFactory.kt
@@ -41,6 +41,7 @@ import software.aws.toolkits.jetbrains.core.credentials.CredentialManager
 import software.aws.toolkits.jetbrains.core.credentials.InteractiveCredential
 import software.aws.toolkits.jetbrains.core.credentials.MfaRequiredInteractiveCredentials
 import software.aws.toolkits.jetbrains.core.credentials.PostValidateInteractiveCredential
+import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
 import software.aws.toolkits.jetbrains.core.credentials.RefreshConnectionAction
 import software.aws.toolkits.jetbrains.core.credentials.SsoRequiredInteractiveCredentials
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager
@@ -133,7 +134,7 @@ class ProfileCredentialsIdentifierSso @TestOnly constructor(
                                     scopes = session.scopes.toList()
                                 )
                             )
-                            reauthConnectionIfNeeded(e.project, connection)
+                            reauthConnectionIfNeeded(e.project, connection, reauthSource = ReauthSource.TOOLKIT)
                             RefreshConnectionAction().actionPerformed(e)
                         }
                     }

diff --git a/...jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt b/...jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt
@@ -102,7 +102,6 @@ class DiskCache(
         LOG.debug { "loadClientRegistration for $cacheKey" }
         val inputStream = clientRegistrationCache(cacheKey).tryInputStreamIfExists()
             ?: return null
-
         return loadClientRegistration(inputStream)
     }
 

diff --git a/...munity/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt b/...munity/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt
@@ -432,6 +432,7 @@ class SsoAccessTokenProvider(
                 requestId = requestId,
                 result = Result.Failed
             )
+            getLogger<SsoAccessTokenProvider>().warn("RefreshAccessTokenFailed: ${e.message}")
             throw e
         }
     }

diff --git a/...munity/src/software/aws/toolkits/jetbrains/core/gettingstarted/GettingStartedAuthUtils.kt b/...munity/src/software/aws/toolkits/jetbrains/core/gettingstarted/GettingStartedAuthUtils.kt
@@ -8,6 +8,7 @@ import software.aws.toolkits.core.utils.tryOrNull
 import software.aws.toolkits.jetbrains.core.credentials.LegacyManagedBearerSsoConnection
 import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
 import software.aws.toolkits.jetbrains.core.credentials.ProfileSsoManagedBearerSsoConnection
+import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.loginSso
 import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
@@ -230,7 +231,7 @@ fun reauthenticateWithQ(project: Project) {
     val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
     if (connection !is ManagedBearerSsoConnection) return
     pluginAwareExecuteOnPooledThread {
-        reauthConnectionIfNeeded(project, connection, isReAuth = true)
+        reauthConnectionIfNeeded(project, connection, isReAuth = true, reauthSource = ReauthSource.Q_CHAT)
     }
 }
 

diff --git a/...core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/LoginBrowser.kt b/...core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/LoginBrowser.kt
@@ -27,6 +27,7 @@ import software.aws.toolkits.core.utils.tryOrNull
 import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.Login
+import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeCatalystConnection
@@ -380,7 +381,7 @@ abstract class LoginBrowser(
     protected fun reauth(connection: ToolkitConnection?) {
         if (connection is AwsBearerTokenConnection) {
             loginWithBackgroundContext {
-                reauthConnectionIfNeeded(project, connection, onPendingToken, isReAuth = true)
+                reauthConnectionIfNeeded(project, connection, onPendingToken, isReAuth = true, reauthSource = ReauthSource.LOGIN_BROWSER)
             }
             stopAndClearBrowserOpenTimer()
         }

diff --git a/...ore/src/software/aws/toolkits/jetbrains/core/credentials/ConnectionSettingsMenuBuilder.kt b/...ore/src/software/aws/toolkits/jetbrains/core/credentials/ConnectionSettingsMenuBuilder.kt
@@ -244,7 +244,7 @@ class ConnectionSettingsMenuBuilder private constructor() {
             addAll(
                 object : DumbAwareAction(message("credentials.individual_identity.reconnect")) {
                     override fun actionPerformed(e: AnActionEvent) {
-                        reauthConnectionIfNeeded(e.project, value, isReAuth = true)
+                        reauthConnectionIfNeeded(e.project, value, isReAuth = true, reauthSource = ReauthSource.TOOLKIT)
 
                         e.project?.let { ToolkitConnectionManager.getInstance(it).switchConnection(value) }
                     }

diff --git a/...rc/software/aws/toolkits/jetbrains/core/credentials/sono/CodeCatalystCredentialManager.kt b/...rc/software/aws/toolkits/jetbrains/core/credentials/sono/CodeCatalystCredentialManager.kt
@@ -12,6 +12,7 @@ import software.aws.toolkits.core.telemetry.DefaultMetricEvent
 import software.aws.toolkits.core.utils.tryOrNull
 import software.aws.toolkits.jetbrains.core.AwsResourceCache
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
+import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.logoutFromSsoConnection
 import software.aws.toolkits.jetbrains.core.credentials.maybeReauthProviderIfNeeded
@@ -78,7 +79,7 @@ class CodeCatalystCredentialManager {
     fun promptAuth(): BearerTokenProvider? {
         connection()?.let {
             val tokenProvider = provider(it)
-            val reauthRequired = maybeReauthProviderIfNeeded(project, tokenProvider) {}
+            val reauthRequired = maybeReauthProviderIfNeeded(project, ReauthSource.CODECATALYST, tokenProvider) {}
             if (reauthRequired) {
                 val useCurrentCredentials =
                     computeOnEdt {

diff --git a/.../src/software/aws/toolkits/jetbrains/core/explorer/devToolsTab/nodes/actions/SonoLogin.kt b/.../src/software/aws/toolkits/jetbrains/core/explorer/devToolsTab/nodes/actions/SonoLogin.kt
@@ -7,6 +7,7 @@ import com.intellij.icons.AllIcons
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.project.DumbAwareAction
+import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeCatalystConnection
 import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
@@ -23,7 +24,7 @@ class SonoLogin : DumbAwareAction(AllIcons.Actions.Execute) {
         pluginAwareExecuteOnPooledThread {
             val connectionManager = ToolkitConnectionManager.getInstance(project)
             connectionManager.activeConnectionForFeature(CodeCatalystConnection.getInstance())?.let {
-                reauthConnectionIfNeeded(project, it, isReAuth = true)
+                reauthConnectionIfNeeded(project, it, isReAuth = true, reauthSource = ReauthSource.CODECATALYST)
                 project.refreshDevToolTree()
             } ?: run {
                 runInEdt {

diff --git a/...etbrains-core/src/software/aws/toolkits/jetbrains/core/gettingstarted/ToolkitAuthUtils.kt b/...etbrains-core/src/software/aws/toolkits/jetbrains/core/gettingstarted/ToolkitAuthUtils.kt
@@ -9,6 +9,7 @@ import com.intellij.openapi.ui.Messages
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.ConfigFilesFacade
 import software.aws.toolkits.jetbrains.core.credentials.DefaultConfigFilesFacade
+import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
 import software.aws.toolkits.jetbrains.core.credentials.UserConfigSsoSessionProfile
 import software.aws.toolkits.jetbrains.core.credentials.authAndUpdateConfig
 import software.aws.toolkits.jetbrains.core.credentials.profiles.SsoSessionConstants
@@ -42,7 +43,7 @@ fun rolePopupFromConnection(
                     Messages.showErrorDialog(project, e.message, message("gettingstarted.explorer.iam.add"))
                 } ?: return@runInEdt
             } else {
-                reauthConnectionIfNeeded(project, connection)
+                reauthConnectionIfNeeded(project, connection, reauthSource = ReauthSource.TOOLKIT)
                 connection
             }.getConnectionSettings().tokenProvider
 

diff --git a/...olkit/jetbrains-ultimate/src/software/aws/toolkits/jetbrains/services/rds/auth/IamAuth.kt b/...olkit/jetbrains-ultimate/src/software/aws/toolkits/jetbrains/services/rds/auth/IamAuth.kt
@@ -22,6 +22,7 @@ import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
 import software.aws.toolkits.jetbrains.core.credentials.CredentialManager
+import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager
 import software.aws.toolkits.jetbrains.core.credentials.UserConfigSsoSessionProfile
 import software.aws.toolkits.jetbrains.core.credentials.profiles.ProfileCredentialsIdentifierSso
@@ -117,7 +118,7 @@ class IamAuth : DatabaseAuthProviderCompatabilityAdapter {
             )
         )
 
-        reauthConnectionIfNeeded(project, ssoConnection)
+        reauthConnectionIfNeeded(project, ssoConnection, reauthSource = ReauthSource.TOOLKIT)
         return connection
     }
-Original file line number
+Diff line change
@@ Expand Up / @@ -432,6 +432,7 @@ class SsoAccessTokenProvider( @@
                     requestId = requestId,
                     result = Result.Failed
                 )
+                getLogger<SsoAccessTokenProvider>().warn("RefreshAccessTokenFailed: ${e.message}")
                 throw e
             }
         }
@@ Expand Down @@