Skip to content

Commit f96225d

Browse files
committed
fix(dev): cache tree
1 parent de9d7cf commit f96225d

File tree

11 files changed

+72
-20
lines changed

11 files changed

+72
-20
lines changed

packages/core/src/amazonq/webview/ui/apps/featureDevChatConnector.ts

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -94,12 +94,13 @@ export class Connector {
9494
})
9595
}
9696

97-
onOpenDiff = (tabID: string, filePath: string, deleted: boolean): void => {
97+
onOpenDiff = (tabID: string, filePath: string, deleted: boolean, messageId?: string): void => {
9898
this.sendMessageToExtension({
9999
command: 'open-diff',
100100
tabID,
101101
filePath,
102102
deleted,
103+
messageId,
103104
tabType: 'featuredev',
104105
})
105106
}
@@ -167,7 +168,11 @@ export class Connector {
167168
canBeVoted: true,
168169
codeReference: messageData.references,
169170
// TODO get the backend to store a message id in addition to conversationID
170-
messageId: messageData.messageID ?? messageData.triggerID ?? messageData.conversationID,
171+
messageId:
172+
messageData.codeGenerationId ??
173+
messageData.messageID ??
174+
messageData.triggerID ??
175+
messageData.conversationID,
171176
fileList: {
172177
rootFolderTitle: 'Changes',
173178
filePaths: messageData.filePaths.map((f: DiffTreeFileInfo) => f.zipFilePath),

packages/core/src/amazonq/webview/ui/connector.ts

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,16 @@ export interface CodeReference {
2424
}
2525
}
2626

27+
export interface UploadHistory {
28+
[key: string]: {
29+
uploadId: string
30+
timestamp: number
31+
tabId: string
32+
filePaths: DiffTreeFileInfo[]
33+
deletedFiles: DiffTreeFileInfo[]
34+
}
35+
}
36+
2737
export interface ChatPayload {
2838
chatMessage: string
2939
chatCommand?: string
@@ -355,10 +365,10 @@ export class Connector {
355365
}
356366
}
357367

358-
onOpenDiff = (tabID: string, filePath: string, deleted: boolean): void => {
368+
onOpenDiff = (tabID: string, filePath: string, deleted: boolean, messageId?: string): void => {
359369
switch (this.tabsStorage.getTab(tabID)?.type) {
360370
case 'featuredev':
361-
this.featureDevChatConnector.onOpenDiff(tabID, filePath, deleted)
371+
this.featureDevChatConnector.onOpenDiff(tabID, filePath, deleted, messageId)
362372
break
363373
}
364374
}

packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ type OpenDiffMessage = {
6767
// currently the zip file path
6868
filePath: string
6969
deleted: boolean
70+
codeGenerationId: string
7071
}
7172

7273
type fileClickedMessage = {
@@ -400,7 +401,8 @@ export class FeatureDevController {
400401
deletedFiles,
401402
session.state.references ?? [],
402403
tabID,
403-
session.uploadId
404+
session.uploadId,
405+
session.state.codeGenerationId ?? ''
404406
)
405407

406408
const remainingIterations = session.state.codeGenerationRemainingIterationCount
@@ -686,6 +688,7 @@ export class FeatureDevController {
686688

687689
private async openDiff(message: OpenDiffMessage) {
688690
const tabId: string = message.tabID
691+
const codeGenerationId: string = message.messageId
689692
const zipFilePath: string = message.filePath
690693
const session = await this.sessionStorage.getSession(tabId)
691694
telemetry.amazonq_isReviewedChanges.emit({
@@ -702,7 +705,11 @@ export class FeatureDevController {
702705
const name = path.basename(pathInfos.relativePath)
703706
await openDeletedDiff(pathInfos.absolutePath, name, tabId)
704707
} else {
705-
const rightPath = path.join(session.uploadId, zipFilePath)
708+
let uploadId = session.uploadId
709+
if (session?.state?.uploadHistory && session.state.uploadHistory[codeGenerationId]) {
710+
uploadId = session?.state?.uploadHistory[codeGenerationId].uploadId
711+
}
712+
const rightPath = path.join(uploadId, zipFilePath)
706713
await openDiff(pathInfos.absolutePath, rightPath, tabId)
707714
}
708715
}

packages/core/src/amazonqFeatureDev/controllers/chat/messenger/messenger.ts

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -115,9 +115,12 @@ export class Messenger {
115115
deletedFiles: DeletedFileInfo[],
116116
references: CodeReference[],
117117
tabID: string,
118-
uploadId: string
118+
uploadId: string,
119+
codeGenerationId: string
119120
) {
120-
this.dispatcher.sendCodeResult(new CodeResultMessage(filePaths, deletedFiles, references, tabID, uploadId))
121+
this.dispatcher.sendCodeResult(
122+
new CodeResultMessage(filePaths, deletedFiles, references, tabID, uploadId, codeGenerationId)
123+
)
121124
}
122125

123126
public sendAsyncEventProgress(tabID: string, inProgress: boolean, message: string | undefined) {

packages/core/src/amazonqFeatureDev/session/session.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,7 @@ export class Session {
130130
fs: this.config.fs,
131131
messenger: this.messenger,
132132
telemetry: this.telemetry,
133+
uploadHistory: this.state.uploadHistory,
133134
})
134135

135136
if (resp.nextState) {

packages/core/src/amazonqFeatureDev/session/sessionState.ts

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ import {
3636
import { prepareRepoData } from '../util/files'
3737
import { TelemetryHelper } from '../util/telemetryHelper'
3838
import { uploadCode } from '../util/upload'
39-
import { CodeReference } from '../../amazonq/webview/ui/connector'
39+
import { CodeReference, UploadHistory } from '../../amazonq/webview/ui/connector'
4040
import { isPresent } from '../../shared/utilities/collectionUtils'
4141
import { AuthUtil } from '../../codewhisperer/util/authUtil'
4242
import { randomUUID } from '../../shared/crypto'
@@ -261,6 +261,7 @@ export class CodeGenState extends CodeGenBase implements SessionState {
261261
public references: CodeReference[],
262262
tabID: string,
263263
private currentIteration: number,
264+
public uploadHistory: UploadHistory,
264265
public codeGenerationRemainingIterationCount?: number,
265266
public codeGenerationTotalIterationCount?: number
266267
) {
@@ -308,6 +309,16 @@ export class CodeGenState extends CodeGenBase implements SessionState {
308309
this.codeGenerationRemainingIterationCount = codeGeneration.codeGenerationRemainingIterationCount
309310
this.codeGenerationTotalIterationCount = codeGeneration.codeGenerationTotalIterationCount
310311

312+
if (action.uploadHistory && !action.uploadHistory[codeGenerationId] && codeGenerationId) {
313+
action.uploadHistory[codeGenerationId] = {
314+
timestamp: Date.now(),
315+
uploadId: this.config.uploadId,
316+
filePaths: codeGeneration.newFiles,
317+
deletedFiles: codeGeneration.deletedFiles,
318+
tabId: this.tabID,
319+
}
320+
}
321+
311322
action.telemetry.setAmazonqNumberOfReferences(this.references.length)
312323
action.telemetry.recordUserCodeGenerationTelemetry(span, this.conversationId)
313324
const nextState = new PrepareCodeGenState(
@@ -318,7 +329,9 @@ export class CodeGenState extends CodeGenBase implements SessionState {
318329
this.tabID,
319330
this.currentIteration + 1,
320331
this.codeGenerationRemainingIterationCount,
321-
this.codeGenerationTotalIterationCount
332+
this.codeGenerationTotalIterationCount,
333+
action.uploadHistory,
334+
codeGenerationId
322335
)
323336
return {
324337
nextState,
@@ -338,6 +351,7 @@ export class MockCodeGenState implements SessionState {
338351
public filePaths: NewFileInfo[]
339352
public deletedFiles: DeletedFileInfo[]
340353
public readonly conversationId: string
354+
public readonly codeGenerationId?: string
341355
public readonly uploadId: string
342356

343357
constructor(
@@ -384,7 +398,8 @@ export class MockCodeGenState implements SessionState {
384398
},
385399
],
386400
this.tabID,
387-
this.uploadId
401+
this.uploadId,
402+
this.codeGenerationId ?? ''
388403
)
389404
action.messenger.sendAnswer({
390405
message: undefined,
@@ -431,11 +446,15 @@ export class PrepareCodeGenState implements SessionState {
431446
public tabID: string,
432447
private currentIteration: number,
433448
public codeGenerationRemainingIterationCount?: number,
434-
public codeGenerationTotalIterationCount?: number
449+
public codeGenerationTotalIterationCount?: number,
450+
public uploadHistory: UploadHistory = {},
451+
public codeGenerationId?: string
435452
) {
436453
this.tokenSource = new vscode.CancellationTokenSource()
437454
this.uploadId = config.uploadId
438455
this.conversationId = config.conversationId
456+
this.uploadHistory = uploadHistory
457+
this.codeGenerationId = codeGenerationId
439458
}
440459

441460
updateWorkspaceRoot(workspaceRoot: string) {
@@ -490,7 +509,8 @@ export class PrepareCodeGenState implements SessionState {
490509
this.deletedFiles,
491510
this.references,
492511
this.tabID,
493-
this.currentIteration
512+
this.currentIteration,
513+
this.uploadHistory
494514
)
495515
return nextState.interact(action)
496516
}

packages/core/src/amazonqFeatureDev/types.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import type { CancellationTokenSource } from 'vscode'
99
import { Messenger } from './controllers/chat/messenger/messenger'
1010
import { FeatureDevClient } from './client/featureDev'
1111
import { TelemetryHelper } from './util/telemetryHelper'
12-
import { CodeReference } from '../amazonq/webview/ui/connector'
12+
import { CodeReference, UploadHistory } from '../amazonq/webview/ui/connector'
1313
import { DiffTreeFileInfo } from '../amazonq/webview/ui/diffTree/types'
1414

1515
export type Interaction = {
@@ -61,11 +61,13 @@ export interface SessionState {
6161
readonly phase?: SessionStatePhase
6262
readonly uploadId: string
6363
readonly tokenSource: CancellationTokenSource
64+
readonly codeGenerationId?: string
6465
readonly tabID: string
6566
interact(action: SessionStateAction): Promise<SessionStateInteraction>
6667
updateWorkspaceRoot?: (workspaceRoot: string) => void
6768
codeGenerationRemainingIterationCount?: number
6869
codeGenerationTotalIterationCount?: number
70+
uploadHistory?: UploadHistory
6971
}
7072

7173
export interface SessionStateConfig {
@@ -82,6 +84,7 @@ export interface SessionStateAction {
8284
messenger: Messenger
8385
fs: VirtualFileSystem
8486
telemetry: TelemetryHelper
87+
uploadHistory?: UploadHistory
8588
}
8689

8790
export type NewFileZipContents = { zipFilePath: string; fileContent: string }

packages/core/src/amazonqFeatureDev/views/actions/uiMessageListener.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@ export class UIMessageListener {
108108
tabID: msg.tabID,
109109
filePath: msg.filePath,
110110
deleted: msg.deleted,
111+
messageId: msg.messageId,
111112
})
112113
}
113114

packages/core/src/amazonqFeatureDev/views/connector/connector.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ export class ErrorMessage extends UiMessage {
3333

3434
export class CodeResultMessage extends UiMessage {
3535
readonly message!: string
36+
readonly codeGenerationId!: string
3637
readonly references!: {
3738
information: string
3839
recommendationContentSpan: {
@@ -48,7 +49,8 @@ export class CodeResultMessage extends UiMessage {
4849
readonly deletedFiles: DeletedFileInfo[],
4950
references: CodeReference[],
5051
tabID: string,
51-
conversationID: string
52+
conversationID: string,
53+
codeGenerationId: string
5254
) {
5355
super(tabID)
5456
this.references = references
@@ -64,6 +66,7 @@ export class CodeResultMessage extends UiMessage {
6466
},
6567
}
6668
})
69+
this.codeGenerationId = codeGenerationId
6770
this.conversationID = conversationID
6871
}
6972
}

packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -263,7 +263,7 @@ describe('Controller', () => {
263263
workspaceFolders,
264264
}
265265

266-
const codeGenState = new CodeGenState(testConfig, getFilePaths(controllerSetup), [], [], tabID, 0)
266+
const codeGenState = new CodeGenState(testConfig, getFilePaths(controllerSetup), [], [], tabID, 0, {})
267267
const newSession = await createSession({
268268
messenger: controllerSetup.messenger,
269269
sessionState: codeGenState,

0 commit comments

Comments
 (0)