Skip to content

Commit 8dc4b81

Browse files
tverneytverney
authored
feat(dev): Stop code generation action (#4948)
Currently /dev don't support stop code generation. This PR introduces this functionality (watch the video below). Mynah UI provides onStopChatResponse API which we can hook in the cancellation token provided on VS Code, sharing across an active session, aborting current progress.
1 parent 6a849a8 commit 8dc4b81

File tree

32 files changed

+2816
-2120
lines changed

32 files changed

+2816
-2120
lines changed

‎.changes/next-release/feature-c5c4cc47-126a-4236-82af-70743a805ca8.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "feature",
3+
"description" : "Amazon Q /dev: Add stop generation action"
4+
}

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,8 @@ class FeatureDevApp : AmazonQApp {
4141
"response-body-link-click" to IncomingFeatureDevMessage.ClickedLink::class,
4242
"insert_code_at_cursor_position" to IncomingFeatureDevMessage.InsertCodeAtCursorPosition::class,
4343
"open-diff" to IncomingFeatureDevMessage.OpenDiff::class,
44-
"file-click" to IncomingFeatureDevMessage.FileClicked::class
44+
"file-click" to IncomingFeatureDevMessage.FileClicked::class,
45+
"stop-response" to IncomingFeatureDevMessage.StopResponse::class
4546
)
4647

4748
scope.launch {
@@ -82,6 +83,7 @@ class FeatureDevApp : AmazonQApp {
8283
is IncomingFeatureDevMessage.InsertCodeAtCursorPosition -> inboundAppMessagesHandler.processInsertCodeAtCursorPosition(message)
8384
is IncomingFeatureDevMessage.OpenDiff -> inboundAppMessagesHandler.processOpenDiff(message)
8485
is IncomingFeatureDevMessage.FileClicked -> inboundAppMessagesHandler.processFileClicked(message)
86+
is IncomingFeatureDevMessage.StopResponse -> inboundAppMessagesHandler.processStopMessage(message)
8587
}
8688
}
8789

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,4 +17,5 @@ interface InboundAppMessagesHandler {
1717
suspend fun processInsertCodeAtCursorPosition(message: IncomingFeatureDevMessage.InsertCodeAtCursorPosition)
1818
suspend fun processOpenDiff(message: IncomingFeatureDevMessage.OpenDiff)
1919
suspend fun processFileClicked(message: IncomingFeatureDevMessage.FileClicked)
20+
suspend fun processStopMessage(message: IncomingFeatureDevMessage.StopResponse)
2021
}

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClient.kt‎

Lines changed: 98 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -39,107 +39,135 @@ import java.time.Instant
3939
import software.amazon.awssdk.services.codewhispererruntime.model.ChatTriggerType as SyncChatTriggerType
4040

4141
@Service(Service.Level.PROJECT)
42-
class FeatureDevClient(private val project: Project) {
42+
class FeatureDevClient(
43+
private val project: Project,
44+
) {
4345
fun getTelemetryOptOutPreference() =
4446
if (AwsSettings.getInstance().isTelemetryEnabled) {
4547
OptOutPreference.OPTIN
4648
} else {
4749
OptOutPreference.OPTOUT
4850
}
4951

50-
private val featureDevUserContext = ClientMetadata.getDefault().let {
51-
val osForFeatureDev: OperatingSystem =
52-
when {
53-
SystemInfo.isWindows -> OperatingSystem.WINDOWS
54-
SystemInfo.isMac -> OperatingSystem.MAC
55-
// For now, categorize everything else as "Linux" (Linux/FreeBSD/Solaris/etc.)
56-
else -> OperatingSystem.LINUX
57-
}
52+
private val featureDevUserContext =
53+
ClientMetadata.getDefault().let {
54+
val osForFeatureDev: OperatingSystem =
55+
when {
56+
SystemInfo.isWindows -> OperatingSystem.WINDOWS
57+
SystemInfo.isMac -> OperatingSystem.MAC
58+
// For now, categorize everything else as "Linux" (Linux/FreeBSD/Solaris/etc.)
59+
else -> OperatingSystem.LINUX
60+
}
5861

59-
UserContext.builder()
60-
.ideCategory(IdeCategory.JETBRAINS)
61-
.operatingSystem(osForFeatureDev)
62-
.product(FEATURE_EVALUATION_PRODUCT_NAME)
63-
.clientId(it.clientId)
64-
.ideVersion(it.awsVersion)
65-
.build()
66-
}
62+
UserContext
63+
.builder()
64+
.ideCategory(IdeCategory.JETBRAINS)
65+
.operatingSystem(osForFeatureDev)
66+
.product(FEATURE_EVALUATION_PRODUCT_NAME)
67+
.clientId(it.clientId)
68+
.ideVersion(it.awsVersion)
69+
.build()
70+
}
6771

68-
private fun connection() = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
69-
?: error("Attempted to use connection while one does not exist")
72+
private fun connection() =
73+
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
74+
?: error("Attempted to use connection while one does not exist")
7075

7176
private fun bearerClient() = connection().getConnectionSettings().awsClient<CodeWhispererRuntimeClient>()
7277

7378
private val amazonQStreamingClient
7479
get() = AmazonQStreamingClient.getInstance(project)
7580

76-
fun sendFeatureDevTelemetryEvent(conversationId: String): SendTelemetryEventResponse = bearerClient().sendTelemetryEvent { requestBuilder ->
77-
requestBuilder.telemetryEvent { telemetryEventBuilder ->
78-
telemetryEventBuilder.featureDevEvent {
79-
it.conversationId(conversationId)
81+
fun sendFeatureDevTelemetryEvent(conversationId: String): SendTelemetryEventResponse =
82+
bearerClient().sendTelemetryEvent { requestBuilder ->
83+
requestBuilder.telemetryEvent { telemetryEventBuilder ->
84+
telemetryEventBuilder.featureDevEvent {
85+
it.conversationId(conversationId)
86+
}
8087
}
88+
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
89+
requestBuilder.userContext(featureDevUserContext)
8190
}
82-
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
83-
requestBuilder.userContext(featureDevUserContext)
84-
}
8591

86-
fun createTaskAssistConversation(): CreateTaskAssistConversationResponse = bearerClient().createTaskAssistConversation(
87-
CreateTaskAssistConversationRequest.builder().build()
88-
)
89-
90-
fun createTaskAssistUploadUrl(conversationId: String, contentChecksumSha256: String, contentLength: Long): CreateUploadUrlResponse =
92+
fun createTaskAssistConversation(): CreateTaskAssistConversationResponse =
93+
bearerClient().createTaskAssistConversation(
94+
CreateTaskAssistConversationRequest.builder().build(),
95+
)
96+
97+
fun createTaskAssistUploadUrl(
98+
conversationId: String,
99+
contentChecksumSha256: String,
100+
contentLength: Long,
101+
uploadId: String,
102+
): CreateUploadUrlResponse =
91103
bearerClient().createUploadUrl {
92-
it.contentChecksumType(ContentChecksumType.SHA_256)
104+
it
105+
.contentChecksumType(ContentChecksumType.SHA_256)
106+
.uploadId(uploadId)
93107
.contentChecksum(contentChecksumSha256)
94108
.contentLength(contentLength)
95109
.artifactType(ArtifactType.SOURCE_CODE)
96110
.uploadIntent(UploadIntent.TASK_ASSIST_PLANNING)
97111
.uploadContext(
98-
UploadContext.builder()
112+
UploadContext
113+
.builder()
99114
.taskAssistPlanningUploadContext(
100-
TaskAssistPlanningUploadContext.builder()
115+
TaskAssistPlanningUploadContext
116+
.builder()
101117
.conversationId(conversationId)
102-
.build()
103-
)
104-
.build()
118+
.build(),
119+
).build(),
105120
)
106121
}
107122

108-
fun startTaskAssistCodeGeneration(conversationId: String, uploadId: String, userMessage: String): StartTaskAssistCodeGenerationResponse = bearerClient()
109-
.startTaskAssistCodeGeneration {
110-
request ->
111-
request
112-
.conversationState {
113-
it
114-
.conversationId(conversationId)
115-
.chatTriggerType(SyncChatTriggerType.MANUAL)
116-
.currentMessage { cm -> cm.userInputMessage { um -> um.content(userMessage) } }
117-
}
118-
.workspaceState {
119-
it
120-
.programmingLanguage { pl -> pl.languageName("javascript") } // This parameter is omitted by featureDev but required in the request
121-
.uploadId(uploadId)
122-
}
123-
}
123+
fun startTaskAssistCodeGeneration(
124+
conversationId: String,
125+
uploadId: String,
126+
userMessage: String,
127+
codeGenerationId: String?,
128+
currentCodeGenerationId: String?,
129+
): StartTaskAssistCodeGenerationResponse =
130+
bearerClient()
131+
.startTaskAssistCodeGeneration { request ->
132+
request
133+
.conversationState {
134+
it
135+
.conversationId(conversationId)
136+
.chatTriggerType(SyncChatTriggerType.MANUAL)
137+
.currentMessage { cm -> cm.userInputMessage { um -> um.content(userMessage) } }
138+
}.workspaceState {
139+
it
140+
.programmingLanguage { pl -> pl.languageName("javascript") } // This parameter is omitted by featureDev but required in the request
141+
.uploadId(uploadId)
142+
}.codeGenerationId(codeGenerationId.toString())
143+
.currentCodeGenerationId(currentCodeGenerationId)
144+
}
124145

125-
fun getTaskAssistCodeGeneration(conversationId: String, codeGenerationId: String): GetTaskAssistCodeGenerationResponse = bearerClient()
126-
.getTaskAssistCodeGeneration {
127-
it
128-
.conversationId(conversationId)
129-
.codeGenerationId(codeGenerationId)
130-
}
146+
fun getTaskAssistCodeGeneration(
147+
conversationId: String,
148+
codeGenerationId: String,
149+
): GetTaskAssistCodeGenerationResponse =
150+
bearerClient()
151+
.getTaskAssistCodeGeneration {
152+
it
153+
.conversationId(conversationId)
154+
.codeGenerationId(codeGenerationId)
155+
}
131156

132-
suspend fun exportTaskAssistResultArchive(conversationId: String): MutableList<ByteArray> = amazonQStreamingClient.exportResultArchive(
133-
conversationId,
134-
ExportIntent.TASK_ASSIST,
135-
null,
136-
{ e ->
137-
LOG.error(e) { "TaskAssist - ExportResultArchive stream exportId=$conversationId exportIntent=${ExportIntent.TASK_ASSIST} Failed: ${e.message} " }
138-
},
139-
{ startTime ->
140-
LOG.info { "TaskAssist - ExportResultArchive latency: ${calculateTotalLatency(startTime, Instant.now())}" }
141-
}
142-
)
157+
suspend fun exportTaskAssistResultArchive(conversationId: String): MutableList<ByteArray> =
158+
amazonQStreamingClient.exportResultArchive(
159+
conversationId,
160+
ExportIntent.TASK_ASSIST,
161+
null,
162+
{ e ->
163+
LOG.error(
164+
e,
165+
) { "TaskAssist - ExportResultArchive stream exportId=$conversationId exportIntent=${ExportIntent.TASK_ASSIST} Failed: ${e.message} " }
166+
},
167+
{ startTime ->
168+
LOG.info { "TaskAssist - ExportResultArchive latency: ${calculateTotalLatency(startTime, Instant.now())}" }
169+
},
170+
)
143171

144172
companion object {
145173
private val LOG = getLogger<FeatureDevClient>()

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import com.intellij.openapi.command.WriteCommandAction
1515
import com.intellij.openapi.editor.Caret
1616
import com.intellij.openapi.editor.Editor
1717
import com.intellij.openapi.fileEditor.FileEditorManager
18+
import com.intellij.openapi.project.Project
1819
import com.intellij.openapi.vfs.VfsUtil
1920
import com.intellij.openapi.wm.ToolWindowManager
2021
import kotlinx.coroutines.withContext
@@ -74,6 +75,7 @@ import software.aws.toolkits.jetbrains.utils.notifyError
7475
import software.aws.toolkits.resources.message
7576
import software.aws.toolkits.telemetry.AmazonqTelemetry
7677
import software.aws.toolkits.telemetry.Result
78+
import software.aws.toolkits.telemetry.UiTelemetry
7779
import java.util.UUID
7880

7981
class FeatureDevController(
@@ -92,6 +94,10 @@ class FeatureDevController(
9294
)
9395
}
9496

97+
override suspend fun processStopMessage(message: IncomingFeatureDevMessage.StopResponse) {
98+
handleStopMessage(message)
99+
}
100+
95101
override suspend fun processNewTabCreatedMessage(message: IncomingFeatureDevMessage.NewTabCreated) {
96102
newTabOpened(message.tabId)
97103
}
@@ -284,6 +290,26 @@ class FeatureDevController(
284290
}
285291
}
286292

293+
private suspend fun handleStopMessage(message: IncomingFeatureDevMessage.StopResponse) {
294+
val session: Session?
295+
UiTelemetry.click(null as Project?, "amazonq_stopCodeGeneration")
296+
messenger.sendAnswer(
297+
tabId = message.tabId,
298+
message("amazonqFeatureDev.code_generation.stopping_code_generation"),
299+
messageType = FeatureDevMessageType.Answer,
300+
canBeVoted = false
301+
)
302+
messenger.sendUpdatePlaceholder(
303+
tabId = message.tabId,
304+
newPlaceholder = message("amazonqFeatureDev.code_generation.stopping_code_generation")
305+
)
306+
messenger.sendChatInputEnabledMessage(tabId = message.tabId, enabled = false)
307+
session = getSessionInfo(message.tabId)
308+
309+
if (session.sessionState.token?.token !== null) {
310+
session.sessionState.token?.cancel()
311+
}
312+
}
287313
private suspend fun insertCode(tabId: String) {
288314
var session: Session? = null
289315
try {

0 commit comments

Comments
 (0)