Add source to load registration metrics

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

@@ -101,15 +101,15 @@ class DiskCache(
         clientRegistrationCache(ssoRegion).tryDeleteIfExists()
     }
 
-    override fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey): ClientRegistration? {
-        LOG.info { "loadClientRegistration for $cacheKey" }
+    override fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey, source: String): ClientRegistration? {
+        LOG.info { "loadClientRegistration:$source for $cacheKey" }
         val inputStream = clientRegistrationCache(cacheKey).tryInputStreamIfExists()
         if (inputStream == null) {
             val stage = LoadCredentialStage.ACCESS_FILE
             LOG.info { "Failed to load Client Registration: cache file does not exist" }
             AuthTelemetry.modifyConnection(
                 action = "Load cache file",
-                source = "loadClientRegistration",
+                source = "loadClientRegistration:$source",
                 result = Result.Failed,
                 reason = "Failed to load Client Registration",
                 reasonDesc = "Load Step:$stage failed. Cache file does not exist"

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

@@ -194,7 +194,7 @@
 
     @Deprecated("Device authorization grant flow is deprecated")
     private fun registerDAGClient(): ClientRegistration {
-        loadDagClientRegistration()?.let {
+        loadDagClientRegistration(SourceOfLoadRegistration.REGISTER_CLIENT.toString())?.let {
             return it
         }
 
@@ -235,7 +235,7 @@
     }
 
     private fun registerPkceClient(): PKCEClientRegistration {
-        loadPkceClientRegistration()?.let {
+        loadPkceClientRegistration(SourceOfLoadRegistration.REGISTER_CLIENT.toString())?.let {
             return it
         }
 
@@ -431,8 +431,8 @@
         stageName = RefreshCredentialStage.LOAD_REGISTRATION
         val registration = try {
             when (currentToken) {
-                is DeviceAuthorizationGrantToken -> loadDagClientRegistration()
-                is PKCEAuthorizationGrantToken -> loadPkceClientRegistration()
+                is DeviceAuthorizationGrantToken -> loadDagClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString())
+                is PKCEAuthorizationGrantToken -> loadPkceClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString())
             }
         } catch (e: Exception) {
             val message = e.message ?: "$stageName: ${e::class.java.name}"
@@ -505,6 +505,11 @@
         }
     }
 
+    enum class SourceOfLoadRegistration {
+        REGISTER_CLIENT,
+        REFRESH_TOKEN,
+    }
+
     private enum class RefreshCredentialStage {
         VALIDATE_REFRESH_TOKEN,
         LOAD_REGISTRATION,
@@ -514,13 +519,13 @@
         SAVE_TOKEN,
     }
 
-    private fun loadDagClientRegistration(): ClientRegistration? =
-        cache.loadClientRegistration(dagClientRegistrationCacheKey)?.let {
+    private fun loadDagClientRegistration(source: String): ClientRegistration? =
+        cache.loadClientRegistration(dagClientRegistrationCacheKey, source)?.let {
             return it
         }
 
-    private fun loadPkceClientRegistration(): PKCEClientRegistration? =
-        cache.loadClientRegistration(pkceClientRegistrationCacheKey)?.let {
+    private fun loadPkceClientRegistration(source: String): PKCEClientRegistration? =
+        cache.loadClientRegistration(pkceClientRegistrationCacheKey, source)?.let {
             return it as PKCEClientRegistration
         }
 

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

@@ -7,7 +7,7 @@ interface SsoCache {
     fun invalidateClientRegistration(ssoRegion: String)
     fun invalidateAccessToken(ssoUrl: String)
 
-    fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey): ClientRegistration?
+    fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey, source: String): ClientRegistration?
     fun saveClientRegistration(cacheKey: ClientRegistrationCacheKey, registration: ClientRegistration)
     fun invalidateClientRegistration(cacheKey: ClientRegistrationCacheKey)
 

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

@@ -57,7 +57,8 @@ class DiskCacheTest {
                     startUrl = ssoUrl,
                     scopes = scopes,
                     region = ssoRegion
-                )
+                ),
+                "testSource"
             )
         ).isNull()
     }
@@ -71,7 +72,7 @@ class DiskCacheTest {
         )
         cacheLocation.resolve("223224b6f0b4702c1a984be8284fe2c9d9718759.json").writeText("badData")
 
-        assertThat(sut.loadClientRegistration(key)).isNull()
+        assertThat(sut.loadClientRegistration(key, "testSource")).isNull()
     }
 
     @Test
@@ -91,7 +92,7 @@ class DiskCacheTest {
             """.trimIndent()
         )
 
-        assertThat(sut.loadClientRegistration(key)).isNull()
+        assertThat(sut.loadClientRegistration(key, "testSource")).isNull()
     }
 
     @Test
@@ -112,7 +113,7 @@ class DiskCacheTest {
             """.trimIndent()
         )
 
-        assertThat(sut.loadClientRegistration(key)).isNull()
+        assertThat(sut.loadClientRegistration(key, "testSource")).isNull()
     }
 
     @Test
@@ -134,7 +135,7 @@ class DiskCacheTest {
             """.trimIndent()
         )
 
-        assertThat(sut.loadClientRegistration(key))
+        assertThat(sut.loadClientRegistration(key, "testSource"))
             .usingRecursiveComparison()
             .isEqualTo(
                 DeviceAuthorizationClientRegistration(
@@ -217,7 +218,7 @@ class DiskCacheTest {
                 """.trimIndent()
             )
 
-        assertThat(sut.loadClientRegistration(key))
+        assertThat(sut.loadClientRegistration(key, "testSource"))
             .usingRecursiveComparison()
             .isEqualTo(
                 PKCEClientRegistration(
@@ -323,10 +324,10 @@ class DiskCacheTest {
             )
         )
 
-        assertThat(sut.loadClientRegistration(key1))
+        assertThat(sut.loadClientRegistration(key1, "testSource"))
             .usingRecursiveComparison()
             .isEqualTo(
-                sut.loadClientRegistration(key2)
+                sut.loadClientRegistration(key2, "testSource")
             )
     }
 
@@ -350,11 +351,11 @@ class DiskCacheTest {
             region = ssoRegion
         )
 
-        assertThat(sut.loadClientRegistration(key)).isNotNull()
+        assertThat(sut.loadClientRegistration(key, "testSource")).isNotNull()
 
         sut.invalidateClientRegistration(key)
 
-        assertThat(sut.loadClientRegistration(key)).isNull()
+        assertThat(sut.loadClientRegistration(key, "testSource")).isNull()
         assertThat(cacheFile).doesNotExist()
     }
 
@@ -619,7 +620,7 @@ class DiskCacheTest {
         registration.setPosixFilePermissions(emptySet())
         assertPosixPermissions(registration, "---------")
 
-        assertThat(sut.loadClientRegistration(key)).isNotNull()
+        assertThat(sut.loadClientRegistration(key, "testSource")).isNotNull()
 
         assertPosixPermissions(registration, "rw-------")
     }

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

@@ -124,7 +124,7 @@ class SsoAccessTokenProviderTest {
         verify(ssoOidcClient).startDeviceAuthorization(any<StartDeviceAuthorizationRequest>())
         verify(ssoOidcClient).createToken(any<CreateTokenRequest>())
         verify(ssoCache).loadAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl })
-        verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion })
+        verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any<String>())
         verify(ssoCache).saveAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl }, eq(accessToken))
     }
 
@@ -170,7 +170,7 @@ class SsoAccessTokenProviderTest {
         verify(ssoOidcClient).startDeviceAuthorization(any<StartDeviceAuthorizationRequest>())
         verify(ssoOidcClient).createToken(any<CreateTokenRequest>())
         verify(ssoCache).loadAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl })
-        verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion })
+        verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion }, any<String>())
         verify(ssoCache).saveClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion }, any())
         verify(ssoCache).saveAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl }, eq(accessToken))
     }
@@ -267,7 +267,7 @@ class SsoAccessTokenProviderTest {
         verify(ssoOidcClient).startDeviceAuthorization(any<StartDeviceAuthorizationRequest>())
         verify(ssoOidcClient, times(2)).createToken(any<CreateTokenRequest>())
         verify(ssoCache).loadAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl })
-        verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion })
+        verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion }, any<String>())
         verify(ssoCache).saveAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl }, eq(accessToken))
     }
 
@@ -296,7 +296,7 @@ class SsoAccessTokenProviderTest {
         val refreshedToken = runBlocking { sut.refreshToken(sut.accessToken()) }
 
         verify(ssoCache).loadAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl })
-        verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion })
+        verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any<String>())
         verify(ssoOidcClient).createToken(any<CreateTokenRequest>())
         verify(ssoCache).saveAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl }, eq(refreshedToken))
     }
@@ -342,7 +342,7 @@ class SsoAccessTokenProviderTest {
             )
 
             on(
-                ssoCache.loadClientRegistration(any<PKCEClientRegistrationCacheKey>())
+                ssoCache.loadClientRegistration(any<PKCEClientRegistrationCacheKey>(), any<String>())
             ).thenReturn(
                 PKCEClientRegistration(
                     clientType = "public",
@@ -369,7 +369,7 @@ class SsoAccessTokenProviderTest {
         val refreshedToken = runBlocking { sut.refreshToken(sut.accessToken()) }
 
         verify(ssoCache).loadAccessToken(any<PKCEAccessTokenCacheKey>())
-        verify(ssoCache).loadClientRegistration(any<PKCEClientRegistrationCacheKey>())
+        verify(ssoCache).loadClientRegistration(any<PKCEClientRegistrationCacheKey>(), any<String>())
         verify(ssoOidcClient).createToken(any<CreateTokenRequest>())
         verify(ssoCache).saveAccessToken(any<PKCEAccessTokenCacheKey>(), eq(refreshedToken))
     }
@@ -390,7 +390,7 @@ class SsoAccessTokenProviderTest {
         verify(ssoOidcClient).startDeviceAuthorization(any<StartDeviceAuthorizationRequest>())
         verify(ssoOidcClient).createToken(any<CreateTokenRequest>())
         verify(ssoCache).loadAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl })
-        verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion })
+        verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion }, any<String>())
     }
 
     @Test
@@ -432,7 +432,7 @@ class SsoAccessTokenProviderTest {
         verify(ssoOidcClient).startDeviceAuthorization(any<StartDeviceAuthorizationRequest>())
         verify(ssoOidcClient, times(2)).createToken(any<CreateTokenRequest>())
         verify(ssoCache).loadAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl })
-        verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion })
+        verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion }, any<String>())
         verify(ssoCache).saveAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl }, eq(accessToken))
     }
 
@@ -452,7 +452,7 @@ class SsoAccessTokenProviderTest {
 
         verify(ssoOidcClient).registerClient(any<RegisterClientRequest>())
         verify(ssoCache).loadAccessToken(any())
-        verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion })
+        verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any<String>())
     }
 
     @Test
@@ -492,7 +492,7 @@ class SsoAccessTokenProviderTest {
             )
 
             on(
-                ssoCache.loadClientRegistration(argThat { region == ssoRegion })
+                ssoCache.loadClientRegistration(argThat { region == ssoRegion }, any<String>())
             ).thenReturn(
                 returnValue
             )

plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/bearer/InteractiveBearerTokenProviderTest.kt

@@ -16,6 +16,7 @@ import org.junit.jupiter.api.assertThrows
 import org.mockito.Mockito
 import org.mockito.kotlin.any
 import org.mockito.kotlin.argThat
+import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.spy
 import org.mockito.kotlin.times
@@ -273,7 +274,7 @@ class InteractiveBearerTokenProviderTest {
     )
 
     private fun stubClientRegistration() {
-        whenever(diskCache.loadClientRegistration(any<DeviceAuthorizationClientRegistrationCacheKey>())).thenReturn(
+        whenever(diskCache.loadClientRegistration(any<DeviceAuthorizationClientRegistrationCacheKey>(), eq("testSource"))).thenReturn(
             DeviceAuthorizationClientRegistration(
                 "",
                 "",