@@ -306,6 +306,9 @@ export class CodeGenState extends CodeGenBase implements SessionState {
306
306
this . tokenSource = action . tokenSource
307
307
}
308
308
this . isCancellationRequested = true
309
+ if ( action . tokenSource ) {
310
+ this . tokenSource = action . tokenSource
311
+ }
309
312
action . tokenSource ?. dispose ( )
310
313
action . tokenSource = undefined
311
314
} )
@@ -321,6 +324,9 @@ export class CodeGenState extends CodeGenBase implements SessionState {
321
324
this . currentCodeGenerationId
322
325
)
323
326
327
+ this . currentCodeGenerationId = codeGenerationId
328
+ this . config . currentCodeGenerationId = codeGenerationId
329
+
324
330
if ( ! this . isCancellationRequested ) {
325
331
action . messenger . sendAnswer ( {
326
332
message : i18n ( 'AWS.amazonq.featureDev.pillText.generatingCode' ) ,
@@ -373,10 +379,10 @@ export class CodeGenState extends CodeGenBase implements SessionState {
373
379
this . currentIteration + 1 ,
374
380
this . codeGenerationRemainingIterationCount ,
375
381
this . codeGenerationTotalIterationCount ,
376
- this . currentCodeGenerationId ,
377
382
action . uploadHistory ,
378
383
codeGenerationId ,
379
- this . tokenSource
384
+ this . tokenSource ,
385
+ this . currentCodeGenerationId
380
386
)
381
387
return {
382
388
nextState,
@@ -498,10 +504,10 @@ export class PrepareCodeGenState implements SessionState {
498
504
public currentIteration : number ,
499
505
public codeGenerationRemainingIterationCount ?: number ,
500
506
public codeGenerationTotalIterationCount ?: number ,
501
- public currentCodeGenerationId ?: string ,
502
507
public uploadHistory : UploadHistory = { } ,
503
508
public codeGenerationId ?: string ,
504
- public superTokenSource ?: vscode . CancellationTokenSource
509
+ public superTokenSource ?: vscode . CancellationTokenSource ,
510
+ public currentCodeGenerationId ?: string
505
511
) {
506
512
this . tokenSource = superTokenSource || new vscode . CancellationTokenSource ( )
507
513
this . uploadId = config . uploadId
0 commit comments