Skip to content

Commit f544307

Browse files
Will-ShaoHuaWill-ShaoHua
authored
refactor browser message passing (#4594)
1 parent 3d8a60d commit f544307

File tree

5 files changed

+435
-118
lines changed

5 files changed

+435
-118
lines changed

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

Lines changed: 39 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33

44
package software.aws.toolkits.jetbrains.services.amazonq
55

6-
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
76
import com.intellij.openapi.Disposable
87
import com.intellij.openapi.actionSystem.AnActionEvent
98
import com.intellij.openapi.actionSystem.DataContext
@@ -18,31 +17,27 @@ import com.intellij.ui.dsl.gridLayout.VerticalAlign
1817
import com.intellij.ui.jcef.JBCefApp
1918
import com.intellij.ui.jcef.JBCefJSQuery
2019
import org.cef.CefApp
21-
import software.aws.toolkits.core.utils.debug
2220
import software.aws.toolkits.core.utils.error
2321
import software.aws.toolkits.core.utils.getLogger
2422
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
25-
import software.aws.toolkits.jetbrains.core.credentials.ManagedBearerSsoConnection
2623
import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager
2724
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
2825
import software.aws.toolkits.jetbrains.core.credentials.actions.SsoLogoutAction
2926
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
30-
import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded
3127
import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES
3228
import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
3329
import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
30+
import software.aws.toolkits.jetbrains.core.webview.BrowserMessage
3431
import software.aws.toolkits.jetbrains.core.webview.BrowserState
3532
import software.aws.toolkits.jetbrains.core.webview.LoginBrowser
3633
import software.aws.toolkits.jetbrains.core.webview.WebviewResourceHandlerFactory
3734
import software.aws.toolkits.jetbrains.isDeveloperMode
3835
import software.aws.toolkits.jetbrains.services.amazonq.util.createBrowser
3936
import software.aws.toolkits.jetbrains.utils.isQConnected
4037
import software.aws.toolkits.jetbrains.utils.isQExpired
41-
import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread
4238
import software.aws.toolkits.telemetry.FeatureId
4339
import software.aws.toolkits.telemetry.WebviewTelemetry
4440
import java.awt.event.ActionListener
45-
import java.util.function.Function
4641
import javax.swing.JButton
4742
import javax.swing.JComponent
4843

@@ -104,49 +99,59 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
10499
// TODO: confirm if we need such configuration or the default is fine
105100
override val jcefBrowser = createBrowser(parentDisposable)
106101
private val query = JBCefJSQuery.create(jcefBrowser)
107-
private val objectMapper = jacksonObjectMapper()
108102

109-
private val handler = Function<String, JBCefJSQuery.Response> {
110-
val jsonTree = objectMapper.readTree(it)
111-
val command = jsonTree.get("command").asText()
112-
LOG.debug { "Data received from Q browser: ${jsonTree.asText()}" }
103+
init {
104+
CefApp.getInstance()
105+
.registerSchemeHandlerFactory(
106+
"http",
107+
domain,
108+
WebviewResourceHandlerFactory(
109+
domain = "http://$domain/",
110+
assetUri = "/webview/assets/"
111+
),
112+
)
113113

114-
when (command) {
115-
"prepareUi" -> {
114+
loadWebView(query)
115+
116+
query.addHandler(jcefHandler)
117+
}
118+
119+
fun component(): JComponent? = jcefBrowser.component
120+
121+
override fun handleBrowserMessage(message: BrowserMessage?) {
122+
if (message == null) {
123+
return
124+
}
125+
126+
when (message) {
127+
is BrowserMessage.PrepareUi -> {
116128
this.prepareBrowser(BrowserState(FeatureId.Q, false))
117129
WebviewTelemetry.amazonqSignInOpened(
118130
project,
119131
reAuth = isQExpired(project)
120132
)
121133
}
122134

123-
"selectConnection" -> {
124-
val connId = jsonTree.get("connectionId").asText()
125-
this.selectionSettings[connId]?.let { settings ->
135+
is BrowserMessage.SelectConnection -> {
136+
this.selectionSettings[message.connectionId]?.let { settings ->
126137
settings.onChange(settings.currentSelection)
127138
}
128139
}
129140

130-
"loginBuilderId" -> {
141+
is BrowserMessage.LoginBuilderId -> {
131142
loginBuilderId(Q_SCOPES)
132143
}
133144

134-
"loginIdC" -> {
135-
// TODO: make it type safe, maybe (de)serialize into a data class
136-
val url = jsonTree.get("url").asText()
137-
val region = jsonTree.get("region").asText()
138-
val awsRegion = AwsRegionProvider.getInstance()[region] ?: error("unknown region returned from Q browser")
139-
140-
val scopes = Q_SCOPES
141-
142-
loginIdC(url, awsRegion, scopes)
145+
is BrowserMessage.LoginIdC -> {
146+
val awsRegion = AwsRegionProvider.getInstance()[message.region] ?: error("unknown region returned from Q browser")
147+
loginIdC(message.url, awsRegion, Q_SCOPES)
143148
}
144149

145-
"cancelLogin" -> {
150+
is BrowserMessage.CancelLogin -> {
146151
cancelLogin()
147152
}
148153

149-
"signout" -> {
154+
is BrowserMessage.Signout -> {
150155
(
151156
ToolkitConnectionManager.getInstance(project)
152157
.activeConnectionForFeature(QConnection.getInstance()) as? AwsBearerTokenConnection
@@ -161,42 +166,16 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
161166
}
162167
}
163168

164-
"reauth" -> {
165-
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let { conn ->
166-
if (conn is ManagedBearerSsoConnection) {
167-
pluginAwareExecuteOnPooledThread {
168-
reauthConnectionIfNeeded(project, conn, onPendingToken)
169-
}
170-
}
171-
}
169+
is BrowserMessage.Reauth -> {
170+
reauth(ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()))
172171
}
173172

174-
else -> {
175-
error("received unknown command from Q browser: $command")
173+
is BrowserMessage.LoginIAM, is BrowserMessage.ToggleBrowser -> {
174+
error("QBrowser doesn't support the provided command ${message::class.simpleName}")
176175
}
177176
}
178-
179-
null
180-
}
181-
182-
init {
183-
CefApp.getInstance()
184-
.registerSchemeHandlerFactory(
185-
"http",
186-
domain,
187-
WebviewResourceHandlerFactory(
188-
domain = "http://$domain/",
189-
assetUri = "/webview/assets/"
190-
),
191-
)
192-
193-
loadWebView(query)
194-
195-
query.addHandler(handler)
196177
}
197178

198-
fun component(): JComponent? = jcefBrowser.component
199-
200179
override fun prepareBrowser(state: BrowserState) {
201180
// TODO: duplicate code in ToolkitLoginWebview
202181
selectionSettings.clear()
@@ -230,7 +209,7 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
230209

231210
// available regions
232211
val regions = AwsRegionProvider.getInstance().allRegionsForService("sso").values.let {
233-
objectMapper.writeValueAsString(it)
212+
writeValueAsString(it)
234213
}
235214

236215
// TODO: pass "REAUTH" if connection expires
@@ -251,7 +230,7 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
251230
},
252231
cancellable: ${state.browserCancellable},
253232
feature: '${state.feature}',
254-
existConnections: ${objectMapper.writeValueAsString(selectionSettings.values.map { it.currentSelection }.toList())}
233+
existConnections: ${writeValueAsString(selectionSettings.values.map { it.currentSelection }.toList())}
255234
}
256235
""".trimIndent()
257236
executeJS("window.ideClient.prepareUi($jsonData)")

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/BrowserMessage.kt‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.core.webview
5+
6+
import com.fasterxml.jackson.annotation.JsonSubTypes
7+
import com.fasterxml.jackson.annotation.JsonTypeInfo
8+
9+
/**
10+
* Message received from the login browser
11+
* property name "command", please refer to [AwsLoginBrowser.getWebviewHTML] and Webview package [defs.d.ts]
12+
*/
13+
@JsonTypeInfo(
14+
use = JsonTypeInfo.Id.NAME,
15+
include = JsonTypeInfo.As.PROPERTY,
16+
property = "command"
17+
)
18+
@JsonSubTypes(
19+
JsonSubTypes.Type(value = BrowserMessage.ToggleBrowser::class, name = "toggleBrowser"),
20+
JsonSubTypes.Type(value = BrowserMessage.PrepareUi::class, name = "prepareUi"),
21+
JsonSubTypes.Type(value = BrowserMessage.SelectConnection::class, name = "selectConnection"),
22+
JsonSubTypes.Type(value = BrowserMessage.LoginBuilderId::class, name = "loginBuilderId"),
23+
JsonSubTypes.Type(value = BrowserMessage.LoginIdC::class, name = "loginIdC"),
24+
JsonSubTypes.Type(value = BrowserMessage.LoginIAM::class, name = "loginIAM"),
25+
JsonSubTypes.Type(value = BrowserMessage.CancelLogin::class, name = "cancelLogin"),
26+
JsonSubTypes.Type(value = BrowserMessage.Signout::class, name = "signout"),
27+
JsonSubTypes.Type(value = BrowserMessage.Reauth::class, name = "reauth")
28+
)
29+
sealed interface BrowserMessage {
30+
31+
// FIX_WHEN_MIN_IS_233: data objects are not stable until Kotlin 1.9
32+
// https://kotlinlang.org/docs/whatsnew19.html#stable-data-objects-for-symmetry-with-data-classes
33+
object PrepareUi : BrowserMessage
34+
35+
data class SelectConnection(val connectionId: String) : BrowserMessage
36+
37+
object ToggleBrowser : BrowserMessage
38+
39+
object LoginBuilderId : BrowserMessage
40+
41+
data class LoginIdC(
42+
val url: String,
43+
val region: String,
44+
val feature: String
45+
) : BrowserMessage
46+
47+
data class LoginIAM(
48+
val profileName: String,
49+
val accessKey: String,
50+
val secretKey: String
51+
) : BrowserMessage
52+
53+
object CancelLogin : BrowserMessage
54+
55+
object Signout : BrowserMessage
56+
57+
object Reauth : BrowserMessage
58+
}

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/LoginBrowser.kt‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33

44
package software.aws.toolkits.jetbrains.core.webview
55

6+
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
7+
import com.fasterxml.jackson.module.kotlin.readValue
68
import com.intellij.openapi.application.runInEdt
79
import com.intellij.openapi.progress.ProcessCanceledException
810
import com.intellij.openapi.progress.blockingContext
@@ -13,10 +15,13 @@ import com.intellij.ui.jcef.JBCefBrowserBase
1315
import com.intellij.ui.jcef.JBCefJSQuery
1416
import kotlinx.coroutines.launch
1517
import kotlinx.coroutines.runBlocking
18+
import org.jetbrains.annotations.VisibleForTesting
19+
import org.slf4j.event.Level
1620
import software.aws.toolkits.core.region.AwsRegion
1721
import software.aws.toolkits.core.utils.debug
1822
import software.aws.toolkits.core.utils.error
1923
import software.aws.toolkits.core.utils.getLogger
24+
import software.aws.toolkits.core.utils.tryOrNull
2025
import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
2126
import software.aws.toolkits.jetbrains.core.credentials.AuthProfile
2227
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
@@ -37,6 +42,7 @@ import software.aws.toolkits.telemetry.CredentialType
3742
import software.aws.toolkits.telemetry.FeatureId
3843
import software.aws.toolkits.telemetry.Result
3944
import java.util.concurrent.Future
45+
import java.util.function.Function
4046

4147
data class BrowserState(val feature: FeatureId, val browserCancellable: Boolean = false, val requireReauth: Boolean = false)
4248

@@ -47,8 +53,23 @@ abstract class LoginBrowser(
4753
) {
4854
abstract val jcefBrowser: JBCefBrowserBase
4955

56+
protected val jcefHandler = Function<String, JBCefJSQuery.Response> {
57+
val obj = LOG.tryOrNull("${this::class.simpleName} unable deserialize login browser message: $it", Level.ERROR) {
58+
objectMapper.readValue<BrowserMessage>(it)
59+
}
60+
61+
handleBrowserMessage(obj)
62+
63+
null
64+
}
65+
5066
protected var currentAuthorization: PendingAuthorization? = null
5167

68+
@VisibleForTesting
69+
internal val objectMapper = jacksonObjectMapper()
70+
71+
abstract fun handleBrowserMessage(message: BrowserMessage?)
72+
5273
protected data class BearerConnectionSelectionSettings(val currentSelection: AwsBearerTokenConnection, val onChange: (AwsBearerTokenConnection) -> Unit)
5374

5475
protected val selectionSettings = mutableMapOf<String, BearerConnectionSelectionSettings>()
@@ -71,6 +92,9 @@ abstract class LoginBrowser(
7192
}
7293
}
7394

95+
// TODO: Dumb and will be addressed in followup PRs
96+
protected fun writeValueAsString(o: Any) = objectMapper.writeValueAsString(o)
97+
7498
protected fun cancelLogin() {
7599
// Essentially Authorization becomes a mutable that allows browser and auth to communicate canceled
76100
// status. There might be a risk of race condition here by changing this global, for which effort

0 commit comments

Comments
 (0)