Skip to content

Commit b92fd65

Browse files
leigaolleigaol
committed
merge
2 parents 94cc915 + fdcae23 commit b92fd65

File tree

36 files changed

+886
-106
lines changed

36 files changed

+886
-106
lines changed

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

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ import kotlinx.coroutines.flow.distinctUntilChanged
1818
import kotlinx.coroutines.flow.launchIn
1919
import kotlinx.coroutines.flow.merge
2020
import kotlinx.coroutines.flow.onEach
21-
import kotlinx.coroutines.future.await
2221
import kotlinx.coroutines.launch
2322
import org.cef.browser.CefBrowser
2423
import org.eclipse.lsp4j.Position
@@ -33,6 +32,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
3332
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
3433
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsServerCapabilitiesProvider
3534
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
35+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager.Companion.convertToJsonToSendToChat
3636
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.getTextDocumentIdentifier
3737
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickNotification
3838
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickParams
@@ -62,6 +62,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatN
6262
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatParams
6363
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatPrompt
6464
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatReadyNotification
65+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatUiMessageParams
6566
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ConversationClickRequest
6667
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CopyCodeToClipboardNotification
6768
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CopyCodeToClipboardParams
@@ -92,9 +93,11 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.Promp
9293
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.QuickChatActionRequest
9394
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SEND_CHAT_COMMAND_PROMPT
9495
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_ERROR_MESSAGE
96+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.STOP_CHAT_RESPONSE
9597
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SendChatPromptRequest
9698
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SourceLinkClickNotification
9799
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SourceLinkClickParams
100+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.StopResponseMessage
98101
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabBarActionParams
99102
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabBarActionRequest
100103
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabEventParams
@@ -104,6 +107,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.util.tabType
104107
import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.AmazonQTheme
105108
import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.ThemeBrowserAdapter
106109
import software.aws.toolkits.jetbrains.settings.MeetQSettings
110+
import software.aws.toolkits.resources.AwsCoreBundle
107111
import software.aws.toolkits.telemetry.MetricResult
108112
import software.aws.toolkits.telemetry.Telemetry
109113
import java.util.concurrent.CompletableFuture
@@ -247,6 +251,10 @@ class BrowserConnector(
247251
encryptionManager = this.encryptionManager
248252
encryptionManager?.encrypt(chatParams)?.let { EncryptedChatParams(it, partialResultToken) }?.let { server.sendChatPrompt(it) }
249253
} ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
254+
255+
// We assume there is only one outgoing request per tab because the input is
256+
// blocked when there is an outgoing request
257+
chatCommunicationManager.setInflightRequestForTab(tabId, result)
250258
showResult(result, partialResultToken, tabId, encryptionManager, browser)
251259
}
252260
CHAT_QUICK_ACTION -> {
@@ -264,6 +272,10 @@ class BrowserConnector(
264272
}
265273
} ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
266274

275+
// We assume there is only one outgoing request per tab because the input is
276+
// blocked when there is an outgoing request
277+
chatCommunicationManager.setInflightRequestForTab(tabId, result)
278+
267279
showResult(result, partialResultToken, tabId, encryptionManager, browser)
268280
}
269281
CHAT_LIST_CONVERSATIONS -> {
@@ -334,6 +346,7 @@ class BrowserConnector(
334346
CHAT_TAB_REMOVE -> {
335347
handleChatNotification<TabEventRequest, TabEventParams>(node) { server, params ->
336348
chatCommunicationManager.removePartialChatMessage(params.tabId)
349+
cancelInflightRequests(params.tabId)
337350
server.tabRemove(params)
338351
}
339352
}
@@ -448,6 +461,30 @@ class BrowserConnector(
448461
server.createPrompt(params)
449462
}
450463
}
464+
STOP_CHAT_RESPONSE -> {
465+
val stopResponseRequest = serializer.deserializeChatMessages<StopResponseMessage>(node)
466+
if (!chatCommunicationManager.hasInflightRequest(stopResponseRequest.params.tabId)) {
467+
return
468+
}
469+
cancelInflightRequests(stopResponseRequest.params.tabId)
470+
chatCommunicationManager.removePartialChatMessage(stopResponseRequest.params.tabId)
471+
472+
val paramsJson = Gson().toJson(
473+
// https://github.com/aws/language-servers/blob/1c0d88806087125b6fc561f610cc15e98127c6bf/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts#L403
474+
ChatUiMessageParams(
475+
title = AwsCoreBundle.message("amazonqChat.stopChatResponse"),
476+
body = ""
477+
)
478+
)
479+
480+
val uiMessage = convertToJsonToSendToChat(
481+
command = SEND_CHAT_COMMAND_PROMPT,
482+
tabId = stopResponseRequest.params.tabId,
483+
params = paramsJson.toString(),
484+
isPartialResult = false
485+
)
486+
browser.postChat(uiMessage)
487+
}
451488
}
452489
}
453490

@@ -471,6 +508,7 @@ class BrowserConnector(
471508
isPartialResult = false
472509
)
473510
browser.postChat(messageToChat)
511+
chatCommunicationManager.removeInflightRequestForTab(tabId)
474512
} catch (e: Exception) {
475513
LOG.error { "Failed to send chat message $e" }
476514
chatCommunicationManager.removePartialChatMessage(partialResultToken)
@@ -483,6 +521,14 @@ class BrowserConnector(
483521
)
484522
)
485523
}
524+
525+
}
526+
}
527+
528+
private fun cancelInflightRequests(tabId: String) {
529+
chatCommunicationManager.getInflightRequestForTab(tabId)?.let { request ->
530+
request.cancel(true)
531+
chatCommunicationManager.removeInflightRequestForTab(tabId)
486532
}
487533
}
488534

‎plugins/amazonq/codetransform/jetbrains-community/resources/META-INF/codetransform-ext-java.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
<projectService serviceImplementation="software.aws.toolkits.jetbrains.services.codemodernizer.panels.managers.CodeModernizerBottomWindowPanelManager"/>
1919
<toolWindow id="aws.codewhisperer.codetransform" anchor="bottom" doNotActivateOnStart="true" canCloseContents="false"
2020
factoryClass="software.aws.toolkits.jetbrains.services.codemodernizer.toolwindow.CodeModernizerBottomToolWindowFactory"
21-
icon="AllIcons.Actions.Preview"/>
21+
icon="AllIcons.Actions.Properties"/>
2222
<fileEditorProvider implementation="software.aws.toolkits.jetbrains.services.codemodernizer.plan.CodeModernizerPlanEditorProvider"/>
2323
<fileEditorProvider implementation="software.aws.toolkits.jetbrains.services.codemodernizer.summary.CodeModernizerSummaryEditorProvider"/>
2424
</extensions>

‎plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/ArtifactHandler.kt‎

Lines changed: 3 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ import software.aws.toolkits.core.utils.exists
2424
import software.aws.toolkits.core.utils.getLogger
2525
import software.aws.toolkits.core.utils.info
2626
import software.aws.toolkits.jetbrains.core.coroutines.EDT
27-
import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineBgContext
2827
import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
2928
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.NoTokenInitializedException
3029
import software.aws.toolkits.jetbrains.services.amazonq.CODE_TRANSFORM_TROUBLESHOOT_DOC_DOWNLOAD_ERROR_OVERVIEW
@@ -50,13 +49,13 @@ import software.aws.toolkits.jetbrains.services.codemodernizer.state.CodeModerni
5049
import software.aws.toolkits.jetbrains.services.codemodernizer.utils.getPathToHilArtifactDir
5150
import software.aws.toolkits.jetbrains.services.codemodernizer.utils.isValidCodeTransformConnection
5251
import software.aws.toolkits.jetbrains.services.codemodernizer.utils.openTroubleshootingGuideNotificationAction
52+
import software.aws.toolkits.jetbrains.services.codemodernizer.utils.zipToPath
5353
import software.aws.toolkits.jetbrains.utils.notifyInfo
5454
import software.aws.toolkits.jetbrains.utils.notifyStickyInfo
5555
import software.aws.toolkits.jetbrains.utils.notifyStickyWarn
5656
import software.aws.toolkits.resources.message
5757
import software.aws.toolkits.telemetry.CodeTransformArtifactType
5858
import java.io.File
59-
import java.nio.file.Files
6059
import java.nio.file.Path
6160
import java.time.Instant
6261
import java.util.concurrent.atomic.AtomicBoolean
@@ -112,32 +111,12 @@ class ArtifactHandler(
112111
}
113112
}
114113

115-
suspend fun unzipToPath(byteArrayList: List<ByteArray>, outputDirPath: Path? = null): Pair<Path, Int> {
116-
val zipFilePath = withContext(getCoroutineBgContext()) {
117-
if (outputDirPath == null) {
118-
Files.createTempFile(null, ".zip")
119-
} else {
120-
Files.createTempFile(outputDirPath, null, ".zip")
121-
}
122-
}
123-
var totalDownloadBytes = 0
124-
withContext(getCoroutineBgContext()) {
125-
Files.newOutputStream(zipFilePath).use {
126-
for (bytes in byteArrayList) {
127-
it.write(bytes)
128-
totalDownloadBytes += bytes.size
129-
}
130-
}
131-
}
132-
return zipFilePath to totalDownloadBytes
133-
}
134-
135114
suspend fun downloadHilArtifact(jobId: JobId, artifactId: String, tmpDir: File): CodeTransformHilDownloadArtifact? {
136115
val downloadResultsResponse = clientAdaptor.downloadExportResultArchive(jobId, artifactId)
137116

138117
return try {
139118
val tmpPath = tmpDir.toPath()
140-
val (downloadZipFilePath, _) = unzipToPath(downloadResultsResponse, tmpPath)
119+
val (downloadZipFilePath, _) = zipToPath(downloadResultsResponse, tmpPath)
141120
LOG.info { "Successfully converted the hil artifact download to a zip at ${downloadZipFilePath.toAbsolutePath()}." }
142121
CodeTransformHilDownloadArtifact.create(downloadZipFilePath, getPathToHilArtifactDir(tmpPath))
143122
} catch (e: Exception) {
@@ -211,7 +190,7 @@ class ArtifactHandler(
211190
val totalDownloadBytes: Int
212191
val zipPath: String
213192
try {
214-
val result = unzipToPath(downloadResultsResponse)
193+
val result = zipToPath(downloadResultsResponse)
215194
path = result.first
216195
totalDownloadBytes = result.second
217196
zipPath = path.toAbsolutePath().toString()

‎plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerManager.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -909,7 +909,7 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo
909909
// Add delay between upload complete and trying to resume
910910
delay(500)
911911

912-
codeTransformationSession?.resumeTransformFromHil()
912+
codeTransformationSession?.resumeTransformation()
913913
} else {
914914
throw CodeModernizerException("Cannot create dependency zip for HIL")
915915
}

‎plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerSession.kt‎

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.Transformation
2020
import software.amazon.awssdk.services.codewhispererruntime.model.TransformationProgressUpdateStatus
2121
import software.amazon.awssdk.services.codewhispererruntime.model.TransformationStatus
2222
import software.amazon.awssdk.services.codewhispererruntime.model.TransformationUserActionStatus
23+
import software.amazon.awssdk.services.codewhispererruntime.model.UploadContext
2324
import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationDownloadArtifactType
2425
import software.amazon.awssdk.services.ssooidc.model.SsoOidcException
2526
import software.aws.toolkits.core.utils.Waiters.waitUntil
@@ -174,6 +175,7 @@ class CodeModernizerSession(
174175
}
175176
// for language upgrades, copyResult should always be Successful here, failure cases already handled
176177
val result = sessionContext.createZipWithModuleFiles(copyResult)
178+
sessionContext.originalUploadZipPath = result.payload.toPath()
177179

178180
if (result is ZipCreationResult.Missing1P) {
179181
telemetryErrorMessage = "Missing 1p dependencies"
@@ -283,9 +285,7 @@ class CodeModernizerSession(
283285
return CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.OTHER(e.localizedMessage))
284286
} finally {
285287
telemetry.uploadProject(payloadSize, startTime, true, telemetryErrorMessage)
286-
if (payload != null) {
287-
deleteUploadArtifact(payload)
288-
}
288+
// do not delete upload ZIP; re-used for client-side build
289289
}
290290

291291
// Send upload completion message to chat (only if successful)
@@ -309,12 +309,6 @@ class CodeModernizerSession(
309309
}
310310
}
311311

312-
internal fun deleteUploadArtifact(payload: File) {
313-
if (!payload.delete()) {
314-
LOG.warn { "Unable to delete upload artifact." }
315-
}
316-
}
317-
318312
private fun startJob(uploadId: String): StartTransformationResponse {
319313
val sourceLanguage = sessionContext.sourceJavaVersion.name.toTransformationLanguage()
320314
val targetLanguage = sessionContext.targetJavaVersion.name.toTransformationLanguage()
@@ -341,9 +335,10 @@ class CodeModernizerSession(
341335
*/
342336
fun resumeJob(startTime: Instant, jobId: JobId) = state.putJobHistory(sessionContext, TransformationStatus.STARTED, jobId.id, startTime)
343337

344-
fun resumeTransformFromHil() {
338+
fun resumeTransformation() {
345339
val clientAdaptor = GumbyClient.getInstance(sessionContext.project)
346340
clientAdaptor.resumeCodeTransformation(state.currentJobId as JobId, TransformationUserActionStatus.COMPLETED)
341+
getLogger<CodeModernizerManager>().info { "Successfully resumed transformation with status of COMPLETED" }
347342
}
348343

349344
fun rejectHilAndContinue(): ResumeTransformationResponse {
@@ -390,7 +385,7 @@ class CodeModernizerSession(
390385
/**
391386
* Adapted from [CodeWhispererCodeScanSession]
392387
*/
393-
suspend fun uploadPayload(payload: File): String {
388+
suspend fun uploadPayload(payload: File, uploadContext: UploadContext? = null): String {
394389
val sha256checksum: String = Base64.getEncoder().encodeToString(
395390
withContext(getCoroutineBgContext()) {
396391
DigestUtils.sha256(FileInputStream(payload))
@@ -400,7 +395,7 @@ class CodeModernizerSession(
400395
throw AlreadyDisposedException("Disposed when about to create upload URL")
401396
}
402397
val clientAdaptor = GumbyClient.getInstance(sessionContext.project)
403-
val createUploadUrlResponse = clientAdaptor.createGumbyUploadUrl(sha256checksum)
398+
val createUploadUrlResponse = clientAdaptor.createGumbyUploadUrl(sha256checksum, uploadContext)
404399

405400
LOG.info {
406401
"Uploading project artifact at ${payload.path} with checksum $sha256checksum using uploadId: ${
@@ -428,9 +423,9 @@ class CodeModernizerSession(
428423
createUploadUrlResponse.kmsKeyArn().orEmpty(),
429424
) { shouldStop.get() }
430425
}
431-
LOG.info { "Upload to S3 succeeded" }
426+
LOG.info { "Upload of ${payload.path} to S3 succeeded with upload context of $uploadContext" }
432427
if (!shouldStop.get()) {
433-
LOG.info { "Uploaded artifact. Latency: ${calculateTotalLatency(uploadStartTime, Instant.now())}ms" }
428+
LOG.info { "Uploaded artifact. Latency: ${calculateTotalLatency(uploadStartTime, Instant.now())} ms" }
434429
}
435430
return createUploadUrlResponse.uploadId()
436431
}

0 commit comments

Comments
 (0)