Skip to content

Commit f8ece81

Browse files
MarcoWang3MarcoWang3
committed
fix(amazonq): refactor codes
1 parent 6cd3f62 commit f8ece81

File tree

1 file changed

+89
-115
lines changed

1 file changed

+89
-115
lines changed

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

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

99+
/**
100+
* Common logic for initializing streaming sessions
101+
*/
102+
private async initializeStreamingSession(
103+
streamingChunk: any,
104+
initializingStreamsByFile: Map<string, Set<string>>,
105+
errorPrefix: string
106+
): Promise<void> {
107+
const filePath = streamingChunk.filePath
108+
const isAlreadyInitializing =
109+
filePath &&
110+
initializingStreamsByFile.has(filePath) &&
111+
initializingStreamsByFile.get(filePath)!.has(streamingChunk.toolUseId)
112+
113+
if (!this.isStreamingActive(streamingChunk.toolUseId) && filePath && !isAlreadyInitializing) {
114+
if (!initializingStreamsByFile.has(filePath)) {
115+
initializingStreamsByFile.set(filePath, new Set())
116+
}
117+
initializingStreamsByFile.get(filePath)!.add(streamingChunk.toolUseId)
118+
119+
try {
120+
await this.startStreamingDiffSession(streamingChunk.toolUseId, filePath)
121+
} catch (error) {
122+
getLogger().error(`${errorPrefix} ${streamingChunk.toolUseId}: ${error}`)
123+
if (errorPrefix.includes('fsReplace')) {
124+
// Don't rethrow for fsReplace
125+
} else {
126+
throw error
127+
}
128+
} finally {
129+
if (filePath && initializingStreamsByFile.has(filePath)) {
130+
const toolUseIds = initializingStreamsByFile.get(filePath)!
131+
toolUseIds.delete(streamingChunk.toolUseId)
132+
if (toolUseIds.size === 0) {
133+
initializingStreamsByFile.delete(filePath)
134+
}
135+
}
136+
}
137+
}
138+
}
139+
140+
/**
141+
* Common logic for processing streaming content and cleanup
142+
*/
143+
private async processStreamingContent(
144+
streamingChunk: any,
145+
initializingStreamsByFile: Map<string, Set<string>>
146+
): Promise<void> {
147+
if (streamingChunk.fsWriteParams) {
148+
if (this.streamingDiffController && (this.streamingDiffController as any).updateFsWriteParams) {
149+
;(this.streamingDiffController as any).updateFsWriteParams(
150+
streamingChunk.toolUseId,
151+
streamingChunk.fsWriteParams
152+
)
153+
}
154+
}
155+
156+
await this.streamContentUpdate(
157+
streamingChunk.toolUseId,
158+
streamingChunk.content || '',
159+
streamingChunk.isComplete || false
160+
)
161+
162+
if (!streamingChunk.isComplete || !streamingChunk.filePath) {
163+
return
164+
}
165+
166+
const toolUseIds = initializingStreamsByFile.get(streamingChunk.filePath)
167+
if (!toolUseIds) {
168+
return
169+
}
170+
171+
toolUseIds.delete(streamingChunk.toolUseId)
172+
173+
if (toolUseIds.size === 0) {
174+
initializingStreamsByFile.delete(streamingChunk.filePath)
175+
}
176+
}
177+
99178
/**
100179
* Handle streaming chunk processing for diff animations
101180
*/
@@ -132,64 +211,12 @@ export class DiffAnimationHandler implements vscode.Disposable {
132211

133212
toolChunks.add(chunkHash)
134213

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
214+
await this.initializeStreamingSession(
215+
streamingChunk,
216+
initializingStreamsByFile,
217+
'Failed to initialize fsReplace streaming session for'
177218
)
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-
}
219+
await this.processStreamingContent(streamingChunk, initializingStreamsByFile)
193220
} catch (error) {
194221
getLogger().error(`Failed to process fsReplace streaming chunk: ${error}`)
195222
initializingStreamsByFile.delete(streamingChunk.toolUseId)
@@ -198,65 +225,12 @@ export class DiffAnimationHandler implements vscode.Disposable {
198225
}
199226

200227
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
228+
await this.initializeStreamingSession(
229+
streamingChunk,
230+
initializingStreamsByFile,
231+
'Failed to initialize streaming session for'
244232
)
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-
}
233+
await this.processStreamingContent(streamingChunk, initializingStreamsByFile)
260234
} catch (error) {
261235
getLogger().error(`Failed to process streaming chunk: ${error}`)
262236
}

0 commit comments

Comments
 (0)