Skip to content

Commit 6cd3f62

Browse files
MarcoWang3MarcoWang3
committed
fix(amazonq): refactor the code base
1 parent a170402 commit 6cd3f62

File tree

2 files changed

+172
-162
lines changed

2 files changed

+172
-162
lines changed

‎packages/amazonq/src/lsp/chat/diffAnimation/diffAnimationHandler.ts‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,172 @@ export class DiffAnimationHandler implements vscode.Disposable {
9696
}
9797
}
9898

99+
/**
100+
* Handle streaming chunk processing for diff animations
101+
*/
102+
public async handleStreamingChunk(
103+
streamingChunk: any,
104+
initializingStreamsByFile: Map<string, Set<string>>,
105+
processedChunks: Map<string, Set<string>>
106+
): Promise<void> {
107+
// Handle fsReplace streaming chunks separately
108+
if (streamingChunk.toolName === 'fsReplace') {
109+
try {
110+
const contentHash = streamingChunk.content
111+
? `${streamingChunk.content.substring(0, 50)}-${streamingChunk.content.length}`
112+
: 'empty'
113+
const chunkHash = `${streamingChunk.toolUseId}-${contentHash}-${streamingChunk.fsWriteParams?.pairIndex || 0}-${streamingChunk.isComplete}`
114+
115+
if (!processedChunks.has(streamingChunk.toolUseId)) {
116+
processedChunks.set(streamingChunk.toolUseId, new Set())
117+
}
118+
119+
const toolChunks = processedChunks.get(streamingChunk.toolUseId)!
120+
121+
if (streamingChunk.fsWriteParams?.command === 'fsReplace_diffPair') {
122+
if (toolChunks.has(chunkHash)) {
123+
return
124+
}
125+
} else {
126+
const simpleHash = `${streamingChunk.toolUseId}-${streamingChunk.content?.length || 0}`
127+
if (toolChunks.has(simpleHash) && streamingChunk.isComplete) {
128+
return
129+
}
130+
toolChunks.add(simpleHash)
131+
}
132+
133+
toolChunks.add(chunkHash)
134+
135+
const filePath = streamingChunk.filePath
136+
const isAlreadyInitializing =
137+
filePath &&
138+
initializingStreamsByFile.has(filePath) &&
139+
initializingStreamsByFile.get(filePath)!.has(streamingChunk.toolUseId)
140+
141+
if (!this.isStreamingActive(streamingChunk.toolUseId) && filePath && !isAlreadyInitializing) {
142+
if (!initializingStreamsByFile.has(filePath)) {
143+
initializingStreamsByFile.set(filePath, new Set())
144+
}
145+
initializingStreamsByFile.get(filePath)!.add(streamingChunk.toolUseId)
146+
147+
try {
148+
await this.startStreamingDiffSession(streamingChunk.toolUseId, filePath)
149+
} catch (error) {
150+
getLogger().error(
151+
`Failed to initialize fsReplace streaming session for ${streamingChunk.toolUseId}: ${error}`
152+
)
153+
} finally {
154+
if (filePath && initializingStreamsByFile.has(filePath)) {
155+
const toolUseIds = initializingStreamsByFile.get(filePath)!
156+
toolUseIds.delete(streamingChunk.toolUseId)
157+
if (toolUseIds.size === 0) {
158+
initializingStreamsByFile.delete(filePath)
159+
}
160+
}
161+
}
162+
}
163+
164+
if (streamingChunk.fsWriteParams) {
165+
if (this.streamingDiffController && (this.streamingDiffController as any).updateFsWriteParams) {
166+
;(this.streamingDiffController as any).updateFsWriteParams(
167+
streamingChunk.toolUseId,
168+
streamingChunk.fsWriteParams
169+
)
170+
}
171+
}
172+
173+
await this.streamContentUpdate(
174+
streamingChunk.toolUseId,
175+
streamingChunk.content || '',
176+
streamingChunk.isComplete || false
177+
)
178+
179+
if (!streamingChunk.isComplete || !streamingChunk.filePath) {
180+
return
181+
}
182+
183+
const toolUseIds = initializingStreamsByFile.get(streamingChunk.filePath)
184+
if (!toolUseIds) {
185+
return
186+
}
187+
188+
toolUseIds.delete(streamingChunk.toolUseId)
189+
190+
if (toolUseIds.size === 0) {
191+
initializingStreamsByFile.delete(streamingChunk.filePath)
192+
}
193+
} catch (error) {
194+
getLogger().error(`Failed to process fsReplace streaming chunk: ${error}`)
195+
initializingStreamsByFile.delete(streamingChunk.toolUseId)
196+
}
197+
return
198+
}
199+
200+
try {
201+
const filePath = streamingChunk.filePath
202+
const isAlreadyInitializing =
203+
filePath &&
204+
initializingStreamsByFile.has(filePath) &&
205+
initializingStreamsByFile.get(filePath)!.has(streamingChunk.toolUseId)
206+
207+
if (!this.isStreamingActive(streamingChunk.toolUseId) && filePath && !isAlreadyInitializing) {
208+
if (!initializingStreamsByFile.has(filePath)) {
209+
initializingStreamsByFile.set(filePath, new Set())
210+
}
211+
initializingStreamsByFile.get(filePath)!.add(streamingChunk.toolUseId)
212+
213+
try {
214+
await this.startStreamingDiffSession(streamingChunk.toolUseId, filePath)
215+
} catch (error) {
216+
getLogger().error(
217+
`Failed to initialize streaming session for ${streamingChunk.toolUseId}: ${error}`
218+
)
219+
throw error
220+
} finally {
221+
if (filePath && initializingStreamsByFile.has(filePath)) {
222+
const toolUseIds = initializingStreamsByFile.get(filePath)!
223+
toolUseIds.delete(streamingChunk.toolUseId)
224+
if (toolUseIds.size === 0) {
225+
initializingStreamsByFile.delete(filePath)
226+
}
227+
}
228+
}
229+
}
230+
231+
if (streamingChunk.fsWriteParams) {
232+
if (this.streamingDiffController && (this.streamingDiffController as any).updateFsWriteParams) {
233+
;(this.streamingDiffController as any).updateFsWriteParams(
234+
streamingChunk.toolUseId,
235+
streamingChunk.fsWriteParams
236+
)
237+
}
238+
}
239+
240+
await this.streamContentUpdate(
241+
streamingChunk.toolUseId,
242+
streamingChunk.content || '',
243+
streamingChunk.isComplete || false
244+
)
245+
246+
if (!streamingChunk.isComplete || !streamingChunk.filePath) {
247+
return
248+
}
249+
250+
const toolUseIds = initializingStreamsByFile.get(streamingChunk.filePath)
251+
if (!toolUseIds) {
252+
return
253+
}
254+
255+
toolUseIds.delete(streamingChunk.toolUseId)
256+
257+
if (toolUseIds.size === 0) {
258+
initializingStreamsByFile.delete(streamingChunk.filePath)
259+
}
260+
} catch (error) {
261+
getLogger().error(`Failed to process streaming chunk: ${error}`)
262+
}
263+
}
264+
99265
public async dispose(): Promise<void> {
100266
this.streamingSessions.clear()
101267
this.streamingDiffController.dispose()

‎packages/amazonq/src/lsp/chat/messages.ts‎

Lines changed: 6 additions & 162 deletions
Original file line numberDiff line numberDiff line change
@@ -178,167 +178,6 @@ async function cleanupTempFiles(): Promise<void> {
178178
}
179179
}
180180

181-
/**
182-
* Handle streaming chunk processing for diff animations
183-
*/
184-
async function handleStreamingChunk(
185-
streamingChunk: any,
186-
initializingStreamsByFile: Map<string, Set<string>>,
187-
processedChunks: Map<string, Set<string>>
188-
): Promise<void> {
189-
// Handle fsReplace streaming chunks separately
190-
if (streamingChunk.toolName === 'fsReplace') {
191-
try {
192-
const contentHash = streamingChunk.content
193-
? `${streamingChunk.content.substring(0, 50)}-${streamingChunk.content.length}`
194-
: 'empty'
195-
const chunkHash = `${streamingChunk.toolUseId}-${contentHash}-${streamingChunk.fsWriteParams?.pairIndex || 0}-${streamingChunk.isComplete}`
196-
197-
if (!processedChunks.has(streamingChunk.toolUseId)) {
198-
processedChunks.set(streamingChunk.toolUseId, new Set())
199-
}
200-
201-
const toolChunks = processedChunks.get(streamingChunk.toolUseId)!
202-
203-
if (streamingChunk.fsWriteParams?.command === 'fsReplace_diffPair') {
204-
if (toolChunks.has(chunkHash)) {
205-
return
206-
}
207-
} else {
208-
const simpleHash = `${streamingChunk.toolUseId}-${streamingChunk.content?.length || 0}`
209-
if (toolChunks.has(simpleHash) && streamingChunk.isComplete) {
210-
return
211-
}
212-
toolChunks.add(simpleHash)
213-
}
214-
215-
toolChunks.add(chunkHash)
216-
217-
const animationHandler = getDiffAnimationHandler()
218-
const filePath = streamingChunk.filePath
219-
const isAlreadyInitializing =
220-
filePath &&
221-
initializingStreamsByFile.has(filePath) &&
222-
initializingStreamsByFile.get(filePath)!.has(streamingChunk.toolUseId)
223-
224-
if (!animationHandler.isStreamingActive(streamingChunk.toolUseId) && filePath && !isAlreadyInitializing) {
225-
if (!initializingStreamsByFile.has(filePath)) {
226-
initializingStreamsByFile.set(filePath, new Set())
227-
}
228-
initializingStreamsByFile.get(filePath)!.add(streamingChunk.toolUseId)
229-
230-
try {
231-
await animationHandler.startStreamingDiffSession(streamingChunk.toolUseId, filePath)
232-
} catch (error) {
233-
getLogger().error(
234-
`Failed to initialize fsReplace streaming session for ${streamingChunk.toolUseId}: ${error}`
235-
)
236-
} finally {
237-
if (filePath && initializingStreamsByFile.has(filePath)) {
238-
const toolUseIds = initializingStreamsByFile.get(filePath)!
239-
toolUseIds.delete(streamingChunk.toolUseId)
240-
if (toolUseIds.size === 0) {
241-
initializingStreamsByFile.delete(filePath)
242-
}
243-
}
244-
}
245-
}
246-
247-
if (streamingChunk.fsWriteParams) {
248-
const streamingController = (animationHandler as any).streamingDiffController
249-
if (streamingController && streamingController.updateFsWriteParams) {
250-
streamingController.updateFsWriteParams(streamingChunk.toolUseId, streamingChunk.fsWriteParams)
251-
}
252-
}
253-
254-
await animationHandler.streamContentUpdate(
255-
streamingChunk.toolUseId,
256-
streamingChunk.content || '',
257-
streamingChunk.isComplete || false
258-
)
259-
260-
if (!streamingChunk.isComplete || !streamingChunk.filePath) {
261-
return
262-
}
263-
264-
const toolUseIds = initializingStreamsByFile.get(streamingChunk.filePath)
265-
if (!toolUseIds) {
266-
return
267-
}
268-
269-
toolUseIds.delete(streamingChunk.toolUseId)
270-
271-
if (toolUseIds.size === 0) {
272-
initializingStreamsByFile.delete(streamingChunk.filePath)
273-
}
274-
} catch (error) {
275-
getLogger().error(`Failed to process fsReplace streaming chunk: ${error}`)
276-
initializingStreamsByFile.delete(streamingChunk.toolUseId)
277-
}
278-
return
279-
}
280-
281-
try {
282-
const animationHandler = getDiffAnimationHandler()
283-
const filePath = streamingChunk.filePath
284-
const isAlreadyInitializing =
285-
filePath &&
286-
initializingStreamsByFile.has(filePath) &&
287-
initializingStreamsByFile.get(filePath)!.has(streamingChunk.toolUseId)
288-
289-
if (!animationHandler.isStreamingActive(streamingChunk.toolUseId) && filePath && !isAlreadyInitializing) {
290-
if (!initializingStreamsByFile.has(filePath)) {
291-
initializingStreamsByFile.set(filePath, new Set())
292-
}
293-
initializingStreamsByFile.get(filePath)!.add(streamingChunk.toolUseId)
294-
295-
try {
296-
await animationHandler.startStreamingDiffSession(streamingChunk.toolUseId, filePath)
297-
} catch (error) {
298-
getLogger().error(`Failed to initialize streaming session for ${streamingChunk.toolUseId}: ${error}`)
299-
throw error
300-
} finally {
301-
if (filePath && initializingStreamsByFile.has(filePath)) {
302-
const toolUseIds = initializingStreamsByFile.get(filePath)!
303-
toolUseIds.delete(streamingChunk.toolUseId)
304-
if (toolUseIds.size === 0) {
305-
initializingStreamsByFile.delete(filePath)
306-
}
307-
}
308-
}
309-
}
310-
311-
if (streamingChunk.fsWriteParams) {
312-
const streamingController = (animationHandler as any).streamingDiffController
313-
if (streamingController && streamingController.updateFsWriteParams) {
314-
streamingController.updateFsWriteParams(streamingChunk.toolUseId, streamingChunk.fsWriteParams)
315-
}
316-
}
317-
318-
await animationHandler.streamContentUpdate(
319-
streamingChunk.toolUseId,
320-
streamingChunk.content || '',
321-
streamingChunk.isComplete || false
322-
)
323-
324-
if (!streamingChunk.isComplete || !streamingChunk.filePath) {
325-
return
326-
}
327-
328-
const toolUseIds = initializingStreamsByFile.get(streamingChunk.filePath)
329-
if (!toolUseIds) {
330-
return
331-
}
332-
333-
toolUseIds.delete(streamingChunk.toolUseId)
334-
335-
if (toolUseIds.size === 0) {
336-
initializingStreamsByFile.delete(streamingChunk.filePath)
337-
}
338-
} catch (error) {
339-
getLogger().error(`Failed to process streaming chunk: ${error}`)
340-
}
341-
}
342181
export function registerMessageListeners(
343182
languageClient: LanguageClient,
344183
provider: AmazonQChatViewProvider,
@@ -886,7 +725,12 @@ export function registerMessageListeners(
886725

887726
languageClient.onNotification(chatUpdateNotificationType.method, async (params: ChatUpdateParams) => {
888727
if ((params.data as any)?.streamingChunk) {
889-
await handleStreamingChunk((params.data as any).streamingChunk, initializingStreamsByFile, processedChunks)
728+
const animationHandler = getDiffAnimationHandler()
729+
await animationHandler.handleStreamingChunk(
730+
(params.data as any).streamingChunk,
731+
initializingStreamsByFile,
732+
processedChunks
733+
)
890734
return
891735
}
892736

0 commit comments

Comments
 (0)