Skip to content

Commit e7aed5b

Browse files
author
Kamiel Wanrooij
committed
fix: improve chat rendering if there are additional chat messages
1 parent 294bfec commit e7aed5b

File tree

3 files changed

+78
-71
lines changed

3 files changed

+78
-71
lines changed

‎chat-client/src/client/mynahUi.ts

Lines changed: 42 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -486,39 +486,38 @@ export const createMynahUi = (
486486
header = contextListToHeader(chatResult.contextList)
487487
}
488488

489-
if (chatResult.additionalMessages?.length) {
490-
const store = mynahUi.getTabData(tabId).getStore() || {}
491-
const chatItems = store.chatItems || []
489+
const store = mynahUi.getTabData(tabId).getStore() || {}
490+
const chatItems = store.chatItems || []
492491

492+
if (chatResult.additionalMessages?.length) {
493493
chatResult.additionalMessages.forEach(am => {
494494
const contextHeader = contextListToHeader(am.contextList)
495495
const header = contextHeader || toMynahHeader(am.header) // Is this mutually exclusive?
496496

497497
const chatItem: ChatItem = {
498498
messageId: am.messageId,
499-
body: am.body,
500-
type: ChatItemType.ANSWER,
501-
header:
502-
am.type === 'tool' && am.header?.fileList && am.header.buttons
503-
? {
504-
...header,
505-
fileList: {
506-
...header?.fileList,
507-
fileTreeTitle: '',
508-
hideFileCount: true,
509-
details: toDetailsWithoutIcon(header?.fileList?.details),
510-
},
511-
buttons: header?.buttons?.map(button => ({ ...button, status: 'clear' })),
512-
}
513-
: header,
499+
body: am.type !== 'tool' && (am.body === undefined || am.body === '') ? 'Thinking...' : am.body,
500+
type: am.type === 'tool' ? ChatItemType.ANSWER : ChatItemType.ANSWER_STREAM,
501+
header: am.type === 'tool' && am.header?.fileList && am.header.buttons
502+
? {
503+
...header,
504+
fileList: {
505+
...header?.fileList,
506+
fileTreeTitle: '',
507+
hideFileCount: true,
508+
details: toDetailsWithoutIcon(header?.fileList?.details),
509+
},
510+
buttons: header?.buttons?.map(button => ({ ...button, status: 'clear' })),
511+
}
512+
: header,
514513
buttons: toMynahButtons(am.buttons),
515514

516515
// file diffs in the header need space
517516
fullWidth: am.type === 'tool' && am.header?.fileList ? true : undefined,
518517
padding: am.type === 'tool' && am.header?.fileList ? false : undefined,
519518
}
520-
const message = chatItems.find(ci => ci.messageId === am.messageId)
521-
if (!message) {
519+
520+
if (!chatItems.find(ci => ci.messageId === am.messageId)) {
522521
mynahUi.addChatItem(tabId, chatItem)
523522
} else {
524523
mynahUi.updateChatAnswerWithMessageId(tabId, am.messageId!, chatItem)
@@ -527,12 +526,19 @@ export const createMynahUi = (
527526
}
528527

529528
if (isPartialResult) {
530-
// type for MynahUI differs from ChatResult types so we ignore it
531-
mynahUi.updateLastChatAnswer(tabId, {
532-
...chatResultWithoutType,
529+
const chatItem = {
530+
...chatResult,
531+
body: chatResult.body === undefined || chatResult.body === '' ? 'Thinking...' : chatResult.body,
532+
type: ChatItemType.ANSWER_STREAM,
533533
header: header,
534534
buttons: buttons,
535-
})
535+
}
536+
537+
if (!chatItems.find(ci => ci.messageId === chatResult.messageId)) {
538+
mynahUi.addChatItem(tabId, chatItem)
539+
} else {
540+
mynahUi.updateChatAnswerWithMessageId(tabId, chatResult.messageId!, chatItem)
541+
}
536542
return
537543
}
538544

@@ -574,20 +580,18 @@ export const createMynahUi = (
574580
// we convert 'ANSWER_STREAM' to 'ANSWER' card.
575581
// First, we unset all the properties and then insert all the data as card item type 'ANSWER'.
576582
// It works, because 'addChatResponse' receives aggregated/joined data send in every next progress update.
577-
mynahUi.updateLastChatAnswer(tabId, {
578-
header: undefined,
579-
body: '',
580-
followUp: undefined,
581-
relatedContent: undefined,
582-
canBeVoted: undefined,
583-
codeReference: undefined,
584-
fileList: undefined,
585-
})
586-
587-
mynahUi.endMessageStream(tabId, chatResult.messageId ?? '')
588-
589-
mynahUi.updateLastChatAnswer(tabId, {
590-
type: ChatItemType.ANSWER,
583+
// mynahUi.updateLastChatAnswer(tabId, {
584+
// header: undefined,
585+
// body: '',
586+
// followUp: undefined,
587+
// relatedContent: undefined,
588+
// canBeVoted: undefined,
589+
// codeReference: undefined,
590+
// fileList: undefined,
591+
// })
592+
593+
mynahUi.endMessageStream(tabId, chatResult.messageId ?? '', {
594+
// type: ChatItemType.ANSWER,
591595
header: header,
592596
buttons: buttons,
593597
body: chatResult.body,

‎server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -441,6 +441,7 @@ export class AgenticChatController implements ChatHandlers {
441441
break
442442
default:
443443
await chatResultStream.writeResultBlock({
444+
type: 'tool',
444445
body: `${executeToolMessage(toolUse)}`,
445446
messageId: toolUse.toolUseId,
446447
})
@@ -474,7 +475,11 @@ export class AgenticChatController implements ChatHandlers {
474475
await chatResultStream.writeResultBlock(chatResult)
475476
break
476477
default:
477-
await chatResultStream.writeResultBlock({ body: toolResultMessage(toolUse, result) })
478+
await chatResultStream.writeResultBlock({
479+
type: 'tool',
480+
body: toolResultMessage(toolUse, result),
481+
messageId: toolUse.toolUseId,
482+
})
478483
break
479484
}
480485
} catch (err) {
@@ -1045,6 +1050,7 @@ export class AgenticChatController implements ChatHandlers {
10451050
const requestId = response.$metadata.requestId!
10461051
const chatEventParser = new AgenticChatEventParser(requestId, metric)
10471052
const streamWriter = chatResultStream.getResultStreamWriter()
1053+
10481054
for await (const chatEvent of response.generateAssistantResponseResponse!) {
10491055
const result = chatEventParser.processPartialEvent(chatEvent, contextList)
10501056

@@ -1053,7 +1059,9 @@ export class AgenticChatController implements ChatHandlers {
10531059
return result
10541060
}
10551061

1056-
await streamWriter.write(result.data.chatResult)
1062+
if (chatEvent.assistantResponseEvent) {
1063+
await streamWriter.write(result.data.chatResult)
1064+
}
10571065
}
10581066
await streamWriter.close()
10591067

‎server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.ts

Lines changed: 26 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
import { ChatResult, FileDetails, ChatMessage } from '@aws/language-server-runtimes/protocol'
2+
import { randomUUID } from 'crypto'
23

34
interface ResultStreamWriter {
4-
write(chunk: ChatResult): Promise<void>
5+
write(chunk: ChatResult, final?: boolean): Promise<void>
56
close(): Promise<void>
67
}
78

@@ -20,15 +21,17 @@ export class AgenticChatResultStream {
2021
chatResultBlocks: [] as ChatMessage[],
2122
isLocked: false,
2223
contextFileList: {} as Record<string, FileDetailsWithPath[]>,
24+
uuid: randomUUID(),
25+
messageId: undefined as string | undefined,
2326
}
2427
readonly #sendProgress: (newChatResult: ChatResult | string) => Promise<void>
2528

2629
constructor(sendProgress: (newChatResult: ChatResult | string) => Promise<void>) {
2730
this.#sendProgress = sendProgress
2831
}
2932

30-
getResult(): ChatResult {
31-
return this.#joinResults(this.#state.chatResultBlocks)
33+
getResult(only?: string): ChatResult {
34+
return this.#joinResults(this.#state.chatResultBlocks, only)
3235
}
3336

3437
getContextFileList(toolUseId: string): FileDetailsWithPath[] {
@@ -42,28 +45,22 @@ export class AgenticChatResultStream {
4245
this.#state.contextFileList[toolUseId].push(fileDetails)
4346
}
4447

45-
#joinResults(chatResults: ChatMessage[]): ChatResult {
46-
const tools: Record<string, boolean> = {}
47-
let firstResponseMessageId: string | undefined
48-
49-
for (const result of chatResults) {
50-
if (result.type === 'tool') {
51-
tools[result.messageId || ''] = true
52-
} else if (tools[result.messageId || '']) {
53-
firstResponseMessageId = result.messageId
54-
break
55-
}
56-
}
57-
48+
#joinResults(chatResults: ChatMessage[], only?: string): ChatResult {
5849
const result: ChatResult = {
59-
body: '', // TODO: somehow doesn't stream unless there is content in the primary result message
50+
body: '',
6051
additionalMessages: [],
61-
messageId: firstResponseMessageId,
52+
messageId: this.#state.messageId || this.#state.uuid,
6253
}
6354

64-
return chatResults.reduce<ChatResult>((acc, c) => {
65-
if (c.messageId && c.messageId !== firstResponseMessageId) {
66-
if (acc.additionalMessages!.some(am => am.messageId === c.messageId)) {
55+
return chatResults
56+
.filter(cr => cr.messageId == this.#state.messageId || only === undefined || only === cr.messageId)
57+
.reduce<ChatResult>((acc, c) => {
58+
if (c.messageId === this.#state.messageId) {
59+
return {
60+
...acc,
61+
body: acc.body + AgenticChatResultStream.resultDelimiter + c.body,
62+
}
63+
} else if (acc.additionalMessages!.some(am => am.messageId === c.messageId)) {
6764
return {
6865
...acc,
6966
additionalMessages: acc.additionalMessages!.map(am => ({
@@ -91,18 +88,12 @@ export class AgenticChatResultStream {
9188
additionalMessages: [...acc.additionalMessages!, c],
9289
}
9390
}
94-
} else {
95-
return {
96-
...acc,
97-
body: c.body + AgenticChatResultStream.resultDelimiter + acc.body,
98-
}
99-
}
100-
}, result)
91+
}, result)
10192
}
10293

10394
async writeResultBlock(result: ChatMessage) {
10495
this.#state.chatResultBlocks.push(result)
105-
await this.#sendProgress(this.getResult())
96+
await this.#sendProgress(this.getResult(result.messageId))
10697
}
10798

10899
getResultStreamWriter(): ResultStreamWriter {
@@ -114,8 +105,12 @@ export class AgenticChatResultStream {
114105
let lastResult: ChatResult | undefined
115106

116107
return {
117-
write: async (intermediateChatResult: ChatResult) => {
118-
const combinedResult = this.#joinResults([...this.#state.chatResultBlocks, intermediateChatResult])
108+
write: async (intermediateChatResult: ChatMessage) => {
109+
this.#state.messageId = intermediateChatResult.messageId
110+
const combinedResult = this.#joinResults(
111+
[...this.#state.chatResultBlocks, intermediateChatResult],
112+
intermediateChatResult.messageId
113+
)
119114
lastResult = intermediateChatResult
120115
return await this.#sendProgress(combinedResult)
121116
},

0 commit comments

Comments
 (0)