@@ -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
@@ -183,19 +183,19 @@ class DocController(
183183 }
184184
185185 FollowUpTypes .CANCEL_FOLDER_SELECTION -> {
186- docGenerationTask.folderLevel = DocFolderLevel . ENTIRE_WORKSPACE
186+ docGenerationTask.reset()
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) {
0 commit comments