Skip to content

Commit 1bf3d94

Browse files
KevinDing1KevinDing1
committed
fix(amazonq):fix cross-tab data interference
1 parent 21aba02 commit 1bf3d94

File tree

4 files changed

+62
-35
lines changed

4 files changed

+62
-35
lines changed

‎.changes/next-release/bugfix-79596dfc-37a0-44c8-adf2-a6c87ba806ac.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "bugfix",
3+
"description" : "Amazon Q: Fix data isolation between tabs to prevent interference when using /doc in multiple tabs"
4+
}

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

Lines changed: 27 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -132,9 +132,8 @@ class DocController(
132132
private val authController: AuthController = AuthController(),
133133
) : InboundAppMessagesHandler {
134134
val messenger = context.messagesFromAppToUi
135-
var mode: Mode = Mode.CREATE
136135
val toolWindow = ToolWindowManager.getInstance(context.project).getToolWindow(AmazonQToolWindowFactory.WINDOW_ID)
137-
var docGenerationTask = DocGenerationTask()
136+
private val docGenerationTasks = DocGenerationTasks()
138137

139138
override suspend fun processPromptChatMessage(message: IncomingDocMessage.ChatPrompt) {
140139
handleChat(
@@ -148,7 +147,7 @@ class DocController(
148147
}
149148

150149
override suspend fun processTabRemovedMessage(message: IncomingDocMessage.TabRemoved) {
151-
docGenerationTask.reset()
150+
docGenerationTasks.deleteTask(message.tabId)
152151
chatSessionStorage.deleteSession(message.tabId)
153152
}
154153

@@ -160,6 +159,7 @@ class DocController(
160159

161160
override suspend fun processFollowupClickedMessage(message: IncomingDocMessage.FollowupClicked) {
162161
val session = getSessionInfo(message.tabId)
162+
val docGenerationTask = docGenerationTasks.getTask(message.tabId)
163163

164164
session.preloader(message.followUp.pillText, messenger) // also stores message in session history
165165

@@ -173,7 +173,7 @@ class DocController(
173173
FollowUpTypes.CLOSE_SESSION -> closeSession(message.tabId)
174174
FollowUpTypes.CREATE_DOCUMENTATION -> {
175175
docGenerationTask.interactionType = DocInteractionType.GENERATE_README
176-
mode = Mode.CREATE
176+
docGenerationTask.mode = Mode.CREATE
177177
promptForDocTarget(message.tabId)
178178
}
179179

@@ -187,15 +187,15 @@ class DocController(
187187
newTask(message.tabId)
188188
}
189189

190-
FollowUpTypes.PROCEED_FOLDER_SELECTION -> if (mode == Mode.EDIT) makeChanges(message.tabId) else onDocsGeneration(message)
190+
FollowUpTypes.PROCEED_FOLDER_SELECTION -> if (docGenerationTask.mode == Mode.EDIT) makeChanges(message.tabId) else onDocsGeneration(message)
191191
FollowUpTypes.ACCEPT_CHANGES -> {
192192
docGenerationTask.userDecision = DocUserDecision.ACCEPT
193193
sendDocAcceptanceTelemetry(message.tabId)
194194
acceptChanges(message)
195195
}
196196

197197
FollowUpTypes.MAKE_CHANGES -> {
198-
mode = Mode.EDIT
198+
docGenerationTask.mode = Mode.EDIT
199199
makeChanges(message.tabId)
200200
}
201201

@@ -206,12 +206,12 @@ class DocController(
206206
}
207207

208208
FollowUpTypes.SYNCHRONIZE_DOCUMENTATION -> {
209-
mode = Mode.SYNC
209+
docGenerationTask.mode = Mode.SYNC
210210
promptForDocTarget(message.tabId)
211211
}
212212

213213
FollowUpTypes.EDIT_DOCUMENTATION -> {
214-
mode = Mode.EDIT
214+
docGenerationTask.mode = Mode.EDIT
215215
docGenerationTask.interactionType = DocInteractionType.EDIT_README
216216
promptForDocTarget(message.tabId)
217217
}
@@ -241,7 +241,6 @@ class DocController(
241241
session.sessionState.token?.cancel()
242242
}
243243

244-
docGenerationTask.reset()
245244
newTask(message.tabId)
246245
}
247246

@@ -307,13 +306,14 @@ class DocController(
307306

308307
private suspend fun promptForDocTarget(tabId: String) {
309308
val session = getSessionInfo(tabId)
309+
val docGenerationTask = docGenerationTasks.getTask(tabId)
310310

311311
val currentSourceFolder = session.context.selectedSourceFolder
312312

313313
try {
314314
messenger.sendFolderConfirmationMessage(
315315
tabId = tabId,
316-
message = if (mode == Mode.CREATE) message("amazonqDoc.prompt.create.confirmation") else message("amazonqDoc.prompt.update"),
316+
message = if (docGenerationTask.mode == Mode.CREATE) message("amazonqDoc.prompt.create.confirmation") else message("amazonqDoc.prompt.update"),
317317
folderPath = currentSourceFolder.name,
318318
followUps = listOf(
319319
FollowUp(
@@ -452,6 +452,9 @@ class DocController(
452452
var session: DocSession? = null
453453
try {
454454
session = getSessionInfo(tabId)
455+
val docGenerationTask = docGenerationTasks.getTask(tabId)
456+
docGenerationTask.mode = Mode.NONE
457+
455458
logger.debug { "$FEATURE_NAME: Session created with id: ${session.tabID}" }
456459

457460
val credentialState = authController.getAuthNeededStates(context.project).amazonQ
@@ -528,7 +531,7 @@ class DocController(
528531
}
529532

530533
private suspend fun newTask(tabId: String) {
531-
docGenerationTask = DocGenerationTask()
534+
docGenerationTasks.deleteTask(tabId)
532535
chatSessionStorage.deleteSession(tabId)
533536

534537
messenger.sendAnswer(
@@ -577,7 +580,7 @@ class DocController(
577580
)
578581

579582
messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false)
580-
docGenerationTask.reset()
583+
docGenerationTasks.deleteTask(tabId)
581584
}
582585

583586
private suspend fun provideFeedbackAndRegenerateCode(tabId: String) {
@@ -728,6 +731,7 @@ class DocController(
728731
message: String,
729732
) {
730733
var session: DocSession? = null
734+
val docGenerationTask = docGenerationTasks.getTask(tabId)
731735
try {
732736
logger.debug { "$FEATURE_NAME: Processing message: $message" }
733737
session = getSessionInfo(tabId)
@@ -746,7 +750,7 @@ class DocController(
746750

747751
when (session.sessionState.phase) {
748752
SessionStatePhase.CODEGEN -> {
749-
onCodeGeneration(session, message, tabId, mode)
753+
onCodeGeneration(session, message, tabId, docGenerationTask.mode)
750754
}
751755

752756
else -> null
@@ -756,7 +760,7 @@ class DocController(
756760
is PrepareDocGenerationState -> state.filePaths
757761
else -> emptyList()
758762
}
759-
sendDocGenerationTelemetry(filePaths, session)
763+
sendDocGenerationTelemetry(filePaths, session, docGenerationTask)
760764
broadcastQEvent(QFeatureEvent.INVOCATION)
761765

762766
if (filePaths.isNotEmpty()) {
@@ -767,7 +771,7 @@ class DocController(
767771
} catch (err: Exception) {
768772
// For non edit mode lock the chat input until they explicitly click one of the follow-ups
769773
var isEnableChatInput = false
770-
if (err is DocException && Mode.EDIT == mode) {
774+
if (err is DocException && docGenerationTask.mode == Mode.EDIT) {
771775
isEnableChatInput = err.remainingIterations != null && err.remainingIterations > 0
772776
}
773777

@@ -779,15 +783,16 @@ class DocController(
779783
messenger.sendUpdatePromptProgress(tabId = followUpMessage.tabId, inProgress(progress = 10, message("amazonqDoc.progress_message.scanning")))
780784

781785
val session = getSessionInfo(followUpMessage.tabId)
786+
val docGenerationTask = docGenerationTasks.getTask(followUpMessage.tabId)
782787

783788
messenger.sendAnswer(
784-
message = docGenerationProgressMessage(DocGenerationStep.UPLOAD_TO_S3, this.mode),
789+
message = docGenerationProgressMessage(DocGenerationStep.UPLOAD_TO_S3, docGenerationTask.mode),
785790
messageType = DocMessageType.AnswerPart,
786791
tabId = followUpMessage.tabId,
787792
)
788793

789794
try {
790-
val sessionMessage: String = when (mode) {
795+
val sessionMessage: String = when (docGenerationTask.mode) {
791796
Mode.CREATE -> message("amazonqDoc.session.create")
792797
else -> message("amazonqDoc.session.sync")
793798
}
@@ -821,10 +826,10 @@ class DocController(
821826
return
822827
}
823828

824-
sendDocGenerationTelemetry(filePaths, session)
829+
sendDocGenerationTelemetry(filePaths, session, docGenerationTask)
825830

826831
messenger.sendAnswer(
827-
message = docGenerationProgressMessage(DocGenerationStep.COMPLETE, mode),
832+
message = docGenerationProgressMessage(DocGenerationStep.COMPLETE, docGenerationTask.mode),
828833
messageType = DocMessageType.AnswerPart,
829834
tabId = followUpMessage.tabId,
830835
)
@@ -907,7 +912,6 @@ class DocController(
907912

908913
private suspend fun retryRequests(tabId: String) {
909914
var session: DocSession? = null
910-
docGenerationTask = DocGenerationTask()
911915
try {
912916
messenger.sendAsyncEventProgress(
913917
tabId = tabId,
@@ -954,6 +958,7 @@ class DocController(
954958
val session = getSessionInfo(tabId)
955959
val currentSourceFolder = session.context.selectedSourceFolder
956960
val projectRoot = session.context.projectRoot
961+
val docGenerationTask = docGenerationTasks.getTask(tabId)
957962

958963
withContext(EDT) {
959964
messenger.sendAnswer(
@@ -1017,7 +1022,7 @@ class DocController(
10171022
}
10181023
}
10191024

1020-
private fun sendDocGenerationTelemetry(filePaths: List<NewFileZipInfo>, session: DocSession) {
1025+
private fun sendDocGenerationTelemetry(filePaths: List<NewFileZipInfo>, session: DocSession, docGenerationTask: DocGenerationTask) {
10211026
docGenerationTask.conversationId = session.conversationId
10221027
val (totalGeneratedChars, totalGeneratedLines, totalGeneratedFiles) = session.countedGeneratedContent(filePaths, docGenerationTask.interactionType)
10231028
docGenerationTask.numberOfGeneratedChars = totalGeneratedChars
@@ -1030,6 +1035,7 @@ class DocController(
10301035

10311036
private fun sendDocAcceptanceTelemetry(tabId: String) {
10321037
val session = getSessionInfo(tabId)
1038+
val docGenerationTask = docGenerationTasks.getTask(tabId)
10331039
var filePaths: List<NewFileZipInfo> = emptyList()
10341040

10351041
when (val state = session.sessionState) {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ suspend fun DocController.onCodeGeneration(session: DocSession, message: String,
3131
messenger.sendAsyncEventProgress(tabId, inProgress = true)
3232
messenger.sendUpdatePromptProgress(tabId, inProgress(progress = 10, message("amazonqDoc.progress_message.scanning")))
3333
messenger.sendAnswer(
34-
message = docGenerationProgressMessage(DocGenerationStep.UPLOAD_TO_S3, this.mode),
34+
message = docGenerationProgressMessage(DocGenerationStep.UPLOAD_TO_S3, mode),
3535
messageType = DocMessageType.AnswerPart,
3636
tabId = tabId,
3737
)
@@ -108,7 +108,7 @@ suspend fun DocController.onCodeGeneration(session: DocSession, message: String,
108108
messenger.sendAnswer(
109109
tabId = tabId,
110110
messageType = DocMessageType.Answer,
111-
message = if (this.mode === Mode.CREATE) {
111+
message = if (mode === Mode.CREATE) {
112112
message("amazonqDoc.answer.readmeCreated")
113113
} else {
114114
"${message("amazonqDoc.answer.readmeUpdated")} ${message("amazonqDoc.answer.codeResult")}"

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

Lines changed: 29 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,24 @@ import software.amazon.awssdk.services.codewhispererruntime.model.DocV2Generatio
1111
import software.aws.toolkits.core.utils.debug
1212
import software.aws.toolkits.core.utils.getLogger
1313

14+
class DocGenerationTasks {
15+
private val tasks: MutableMap<String, DocGenerationTask> = mutableMapOf()
16+
17+
fun getTask(tabId: String): DocGenerationTask {
18+
if (!tasks.containsKey(tabId)) {
19+
tasks[tabId] = DocGenerationTask()
20+
}
21+
return tasks[tabId]!!
22+
}
23+
24+
fun deleteTask(tabId: String) {
25+
tasks.remove(tabId)
26+
}
27+
}
28+
1429
class DocGenerationTask {
30+
var mode: Mode = Mode.NONE
31+
1532
// Telemetry fields
1633
var conversationId: String? = null
1734
var numberOfAddedChars: Int? = null
@@ -22,8 +39,8 @@ class DocGenerationTask {
2239
var numberOfGeneratedFiles: Int? = null
2340
var userDecision: DocUserDecision? = null
2441
var interactionType: DocInteractionType? = null
25-
var numberOfNavigations = 0
26-
var folderLevel: DocFolderLevel? = DocFolderLevel.ENTIRE_WORKSPACE
42+
var numberOfNavigations: Int = 0
43+
var folderLevel: DocFolderLevel = DocFolderLevel.ENTIRE_WORKSPACE
2744
fun docGenerationEventBase(): DocV2GenerationEvent {
2845
val undefinedProps = this::class.java.declaredFields
2946
.filter { it.get(this) == null }
@@ -35,10 +52,10 @@ class DocGenerationTask {
3552
}
3653

3754
return DocV2GenerationEvent.builder()
38-
.conversationId(conversationId)
39-
.numberOfGeneratedChars(numberOfGeneratedChars)
40-
.numberOfGeneratedLines(numberOfGeneratedLines)
41-
.numberOfGeneratedFiles(numberOfGeneratedFiles)
55+
.conversationId(conversationId ?: "")
56+
.numberOfGeneratedChars(numberOfGeneratedChars ?: 0)
57+
.numberOfGeneratedLines(numberOfGeneratedLines ?: 0)
58+
.numberOfGeneratedFiles(numberOfGeneratedFiles ?: 0)
4259
.interactionType(interactionType)
4360
.numberOfNavigations(numberOfNavigations)
4461
.folderLevel(folderLevel)
@@ -56,12 +73,12 @@ class DocGenerationTask {
5673
}
5774

5875
return DocV2AcceptanceEvent.builder()
59-
.conversationId(conversationId)
60-
.numberOfAddedChars(numberOfAddedChars)
61-
.numberOfAddedLines(numberOfAddedLines)
62-
.numberOfAddedFiles(numberOfAddedFiles)
63-
.userDecision(userDecision)
64-
.interactionType(interactionType)
76+
.conversationId(conversationId ?: "")
77+
.numberOfAddedChars(numberOfAddedChars ?: 0)
78+
.numberOfAddedLines(numberOfAddedLines ?: 0)
79+
.numberOfAddedFiles(numberOfAddedFiles ?: 0)
80+
.userDecision(userDecision ?: DocUserDecision.ACCEPT)
81+
.interactionType(interactionType ?: DocInteractionType.GENERATE_README)
6582
.numberOfNavigations(numberOfNavigations)
6683
.folderLevel(folderLevel)
6784
.build()

0 commit comments

Comments
 (0)