Skip to content

Commit f8d7360

Browse files
manodnyabmanodnyab
authored
Add metric for location of reauth (#5021)
* Add metric for location of reauth * fixed tests * Removed comment
1 parent 4eb2943 commit f8d7360

File tree

16 files changed

+59
-15
lines changed

16 files changed

+59
-15
lines changed

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/explorerActions/SignInToQAction.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import com.intellij.openapi.actionSystem.AnActionEvent
88
import com.intellij.openapi.application.runInEdt
99
import com.intellij.openapi.project.DumbAwareAction
1010
import com.intellij.openapi.wm.ToolWindowManager
11+
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
1112
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
1213
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
1314
import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
@@ -39,7 +40,7 @@ abstract class SignInToQActionBase(actionName: String) : DumbAwareAction(actionN
3940
UiTelemetry.click(project, "auth_start_Q")
4041
val connectionManager = ToolkitConnectionManager.getInstance(project)
4142
connectionManager.activeConnectionForFeature(QConnection.getInstance())?.let {
42-
reauthConnectionIfNeeded(project, it, isReAuth = true)
43+
reauthConnectionIfNeeded(project, it, isReAuth = true, reauthSource = ReauthSource.CODEWHISPERER_STATUSBAR)
4344
} ?: run {
4445
runInEdt {
4546
if (requestCredentialsForQ(project, isReauth = false)) {

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import software.aws.toolkits.core.utils.getLogger
2626
import software.aws.toolkits.core.utils.warn
2727
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
2828
import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
29+
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
2930
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
3031
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
3132
import software.aws.toolkits.jetbrains.core.credentials.maybeReauthProviderIfNeeded
@@ -171,7 +172,7 @@ object CodeWhispererUtil {
171172
if (!isQExpired(project)) return false
172173
val tokenProvider = tokenProvider(project) ?: return false
173174
return try {
174-
maybeReauthProviderIfNeeded(project, tokenProvider) {
175+
maybeReauthProviderIfNeeded(project, ReauthSource.CODEWHISPERER, tokenProvider) {
175176
runInEdt {
176177
if (!CodeWhispererService.hasReAuthPromptBeenShown()) {
177178
notifyConnectionExpiredRequestReauth(project)
@@ -256,7 +257,7 @@ object CodeWhispererUtil {
256257
val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())
257258
if (connection !is ManagedBearerSsoConnection) return
258259
pluginAwareExecuteOnPooledThread {
259-
reauthConnectionIfNeeded(project, connection, isReAuth = true)
260+
reauthConnectionIfNeeded(project, connection, isReAuth = true, reauthSource = ReauthSource.CODEWHISPERER)
260261
}
261262
}
262263

‎plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtilTest.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import org.junit.jupiter.api.Test
1616
import org.junit.jupiter.api.extension.RegisterExtension
1717
import software.aws.toolkits.core.utils.test.aString
1818
import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
19+
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
1920
import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager
2021
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
2122
import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
@@ -51,7 +52,7 @@ class CodeWhispererUtilTest {
5152
CodeWhispererUtil.reconnectCodeWhisperer(projectExtension.project)
5253

5354
verify {
54-
reauthConnectionIfNeeded(projectExtension.project, mockConnection, isReAuth = true)
55+
reauthConnectionIfNeeded(projectExtension.project, mockConnection, isReAuth = true, reauthSource = ReauthSource.CODEWHISPERER)
5556
}
5657
}
5758
}

‎plugins/core/jetbrains-community/resources/telemetryOverride.json‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,11 @@
273273
"type": "string",
274274
"description": "Comma-delimited list of enabled auths."
275275
},
276+
{
277+
"name": "authRefreshSource",
278+
"type": "string",
279+
"description": "Source triggering token refresh"
280+
},
276281
{
277282
"name": "component",
278283
"allowedValues": [
@@ -813,6 +818,17 @@
813818
"required": true
814819
}
815820
]
821+
},
822+
{
823+
"name": "auth_sourceOfRefresh",
824+
"description": "Source of user triggered refresh",
825+
"metadata": [
826+
{
827+
"type": "authRefreshSource",
828+
"required": false
829+
}
830+
]
816831
}
832+
817833
]
818834
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,7 @@ fun authAndUpdateConfig(
198198

199199
val connection = try {
200200
ToolkitAuthManager.getInstance().tryCreateTransientSsoConnection(updatedProfile) { connection ->
201-
reauthConnectionIfNeeded(project, connection, onPendingToken)
201+
reauthConnectionIfNeeded(project, connection, onPendingToken, reauthSource = ReauthSource.FRESH_AUTH)
202202
}
203203
} catch (e: Exception) {
204204
onError(e)

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

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenPr
2727
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.InteractiveBearerTokenProvider
2828
import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
2929
import software.aws.toolkits.resources.AwsCoreBundle
30+
import software.aws.toolkits.telemetry.AuthTelemetry
3031
import software.aws.toolkits.telemetry.CredentialSourceId
3132
import software.aws.toolkits.telemetry.CredentialType
3233
import software.aws.toolkits.telemetry.Result
@@ -135,6 +136,7 @@ fun loginSso(
135136
onPendingToken = onPendingToken,
136137
isReAuth = false,
137138
source = metadata?.sourceId,
139+
reauthSource = ReauthSource.FRESH_AUTH
138140
)
139141
}
140142
} catch (e: Exception) {
@@ -182,6 +184,7 @@ fun loginSso(
182184
onPendingToken = onPendingToken,
183185
isReAuth = true,
184186
source = metadata?.sourceId,
187+
reauthSource = ReauthSource.COMMON_LOGIN
185188
)
186189
return@let connection
187190
}
@@ -240,6 +243,7 @@ fun reauthConnectionIfNeeded(
240243
onPendingToken: (InteractiveBearerTokenProvider) -> Unit = {},
241244
isReAuth: Boolean = false,
242245
source: String? = null,
246+
reauthSource: ReauthSource? = ReauthSource.TOOLKIT,
243247
): BearerTokenProvider {
244248
val tokenProvider = (connection.getConnectionSettings() as TokenConnectionSettings).tokenProvider.delegate as BearerTokenProvider
245249
if (tokenProvider is InteractiveBearerTokenProvider) {
@@ -248,7 +252,7 @@ fun reauthConnectionIfNeeded(
248252

249253
val startUrl = (connection as AwsBearerTokenConnection).startUrl
250254
var didReauth = false
251-
maybeReauthProviderIfNeeded(project, tokenProvider) {
255+
maybeReauthProviderIfNeeded(project, reauthSource, tokenProvider) {
252256
didReauth = true
253257
runUnderProgressIfNeeded(project, AwsCoreBundle.message("credentials.pending.title"), true) {
254258
try {
@@ -301,6 +305,7 @@ fun reauthConnectionIfNeeded(
301305
// Return true if need to re-auth, false otherwise
302306
fun maybeReauthProviderIfNeeded(
303307
project: Project?,
308+
reauthSource: ReauthSource? = ReauthSource.TOOLKIT,
304309
tokenProvider: BearerTokenProvider,
305310
onReauthRequired: (SsoOidcException?) -> Any,
306311
): Boolean {
@@ -314,12 +319,14 @@ fun maybeReauthProviderIfNeeded(
314319

315320
BearerTokenAuthState.NEEDS_REFRESH -> {
316321
try {
322+
getLogger<ToolkitAuthManager>().warn { "Starting token refresh" }
317323
return runUnderProgressIfNeeded(project, AwsCoreBundle.message("credentials.refreshing"), true) {
318324
tokenProvider.resolveToken()
319325
BearerTokenProviderListener.notifyCredUpdate(tokenProvider.id)
320326
return@runUnderProgressIfNeeded false
321327
}
322328
} catch (e: SsoOidcException) {
329+
AuthTelemetry.sourceOfRefresh(authRefreshSource = reauthSource.toString())
323330
getLogger<ToolkitAuthManager>().warn(e) { "Redriving bearer token login flow since token could not be refreshed" }
324331
onReauthRequired(e)
325332
return true
@@ -332,6 +339,17 @@ fun maybeReauthProviderIfNeeded(
332339
}
333340
}
334341

342+
enum class ReauthSource {
343+
CODEWHISPERER,
344+
TOOLKIT,
345+
Q_CHAT,
346+
LOGIN_BROWSER,
347+
CODEWHISPERER_STATUSBAR,
348+
CODECATALYST,
349+
COMMON_LOGIN,
350+
FRESH_AUTH,
351+
}
352+
335353
fun deleteSsoConnection(connection: ProfileSsoManagedBearerSsoConnection) =
336354
deleteSsoConnection(connection.configSessionName)
337355

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/profiles/ProfileCredentialProviderFactory.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ import software.aws.toolkits.jetbrains.core.credentials.CredentialManager
4141
import software.aws.toolkits.jetbrains.core.credentials.InteractiveCredential
4242
import software.aws.toolkits.jetbrains.core.credentials.MfaRequiredInteractiveCredentials
4343
import software.aws.toolkits.jetbrains.core.credentials.PostValidateInteractiveCredential
44+
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
4445
import software.aws.toolkits.jetbrains.core.credentials.RefreshConnectionAction
4546
import software.aws.toolkits.jetbrains.core.credentials.SsoRequiredInteractiveCredentials
4647
import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager
@@ -133,7 +134,7 @@ class ProfileCredentialsIdentifierSso @TestOnly constructor(
133134
scopes = session.scopes.toList()
134135
)
135136
)
136-
reauthConnectionIfNeeded(e.project, connection)
137+
reauthConnectionIfNeeded(e.project, connection, reauthSource = ReauthSource.TOOLKIT)
137138
RefreshConnectionAction().actionPerformed(e)
138139
}
139140
}

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,6 @@ class DiskCache(
102102
LOG.debug { "loadClientRegistration for $cacheKey" }
103103
val inputStream = clientRegistrationCache(cacheKey).tryInputStreamIfExists()
104104
?: return null
105-
106105
return loadClientRegistration(inputStream)
107106
}
108107

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -432,6 +432,7 @@ class SsoAccessTokenProvider(
432432
requestId = requestId,
433433
result = Result.Failed
434434
)
435+
getLogger<SsoAccessTokenProvider>().warn("RefreshAccessTokenFailed: ${e.message}")
435436
throw e
436437
}
437438
}

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/gettingstarted/GettingStartedAuthUtils.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import software.aws.toolkits.core.utils.tryOrNull
88
import software.aws.toolkits.jetbrains.core.credentials.LegacyManagedBearerSsoConnection
99
import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
1010
import software.aws.toolkits.jetbrains.core.credentials.ProfileSsoManagedBearerSsoConnection
11+
import software.aws.toolkits.jetbrains.core.credentials.ReauthSource
1112
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
1213
import software.aws.toolkits.jetbrains.core.credentials.loginSso
1314
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
@@ -230,7 +231,7 @@ fun reauthenticateWithQ(project: Project) {
230231
val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
231232
if (connection !is ManagedBearerSsoConnection) return
232233
pluginAwareExecuteOnPooledThread {
233-
reauthConnectionIfNeeded(project, connection, isReAuth = true)
234+
reauthConnectionIfNeeded(project, connection, isReAuth = true, reauthSource = ReauthSource.Q_CHAT)
234235
}
235236
}
236237

0 commit comments

Comments
 (0)