Skip to content

Commit 3420883

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

File tree

5 files changed

+88
-88
lines changed

5 files changed

+88
-88
lines changed

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

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ import {
2222
SEND_TO_PROMPT_TELEMETRY_EVENT,
2323
TAB_ADD_TELEMETRY_EVENT,
2424
} from '../contracts/telemetry'
25-
import { ChatItemType, MynahUI } from '@aws/mynah-ui'
25+
import { MynahUI } from '@aws/mynah-ui'
2626
import { TabFactory } from './tabs/tabFactory'
2727
import { ChatClientAdapter } from '../contracts/chatClientAdapter'
2828

@@ -217,10 +217,16 @@ describe('Chat', () => {
217217
})
218218
window.dispatchEvent(chatEvent)
219219

220-
assert.calledOnceWithExactly(endMessageStreamStub, tabId, '')
221-
assert.calledTwice(updateLastChatAnswerStub)
222-
assert.calledWithMatch(updateLastChatAnswerStub, tabId, { body: '' })
223-
assert.calledWithMatch(updateLastChatAnswerStub, tabId, { body, type: ChatItemType.ANSWER })
220+
assert.calledOnceWithExactly(endMessageStreamStub, tabId, '', {
221+
header: { icon: undefined, buttons: undefined },
222+
buttons: undefined,
223+
body: 'some response',
224+
followUp: {},
225+
relatedContent: undefined,
226+
canBeVoted: undefined,
227+
codeReference: undefined,
228+
fileList: undefined,
229+
})
224230
assert.calledOnceWithExactly(updateStoreStub, tabId, {
225231
loadingChat: false,
226232
promptInputDisabledState: false,

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

Lines changed: 19 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -486,18 +486,18 @@ 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,
499+
body: am.type !== 'tool' && (am.body === undefined || am.body === '') ? 'Thinking...' : am.body,
500+
type: am.type === 'tool' ? ChatItemType.ANSWER : ChatItemType.ANSWER_STREAM,
501501
header:
502502
am.type === 'tool' && am.header?.fileList && am.header.buttons
503503
? {
@@ -517,8 +517,8 @@ export const createMynahUi = (
517517
fullWidth: am.type === 'tool' && am.header?.fileList ? true : undefined,
518518
padding: am.type === 'tool' && am.header?.fileList ? false : undefined,
519519
}
520-
const message = chatItems.find(ci => ci.messageId === am.messageId)
521-
if (!message) {
520+
521+
if (!chatItems.find(ci => ci.messageId === am.messageId)) {
522522
mynahUi.addChatItem(tabId, chatItem)
523523
} else {
524524
mynahUi.updateChatAnswerWithMessageId(tabId, am.messageId!, chatItem)
@@ -527,12 +527,19 @@ export const createMynahUi = (
527527
}
528528

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

@@ -569,25 +576,7 @@ export const createMynahUi = (
569576
}
570577
: {}
571578

572-
// TODO: ensure all card item types are supported for export on MynahUI side.
573-
// Chat export does not work with 'ANSWER_STREAM' cards, so at the end of the streaming
574-
// we convert 'ANSWER_STREAM' to 'ANSWER' card.
575-
// First, we unset all the properties and then insert all the data as card item type 'ANSWER'.
576-
// 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,
579+
mynahUi.endMessageStream(tabId, chatResult.messageId ?? '', {
591580
header: header,
592581
buttons: buttons,
593582
body: chatResult.body,

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

Lines changed: 22 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -50,11 +50,6 @@ describe('AgenticChatController', () => {
5050
const mockMessageId = 'mock-message-id'
5151

5252
const mockChatResponseList: ChatResponseStream[] = [
53-
{
54-
messageMetadataEvent: {
55-
conversationId: mockConversationId,
56-
},
57-
},
5853
{
5954
assistantResponseEvent: {
6055
content: 'Hello ',
@@ -73,18 +68,13 @@ describe('AgenticChatController', () => {
7368
]
7469

7570
const expectedCompleteChatResult: ChatResult = {
76-
body: '',
77-
messageId: undefined,
78-
additionalMessages: [
79-
{
80-
body: 'Hello World!',
81-
canBeVoted: true,
82-
messageId: 'mock-message-id',
83-
codeReference: undefined,
84-
followUp: undefined,
85-
relatedContent: undefined,
86-
},
87-
],
71+
body: 'Hello World!',
72+
canBeVoted: true,
73+
messageId: 'mock-message-id',
74+
codeReference: undefined,
75+
followUp: undefined,
76+
relatedContent: undefined,
77+
additionalMessages: [],
8878
}
8979

9080
const expectedCompleteInlineChatResult: InlineChatResult = {
@@ -337,7 +327,11 @@ describe('AgenticChatController', () => {
337327
const chatResult = await chatResultPromise
338328

339329
sinon.assert.callCount(testFeatures.lsp.sendProgress, 0)
340-
assert.deepStrictEqual(chatResult, expectedCompleteChatResult)
330+
assert.deepStrictEqual(chatResult, {
331+
additionalMessages: [],
332+
body: '\n\nHello World!',
333+
messageId: 'mock-message-id',
334+
})
341335
})
342336

343337
it('creates a new conversationId if missing in the session', async () => {
@@ -860,7 +854,11 @@ describe('AgenticChatController', () => {
860854
const chatResult = await chatResultPromise
861855

862856
sinon.assert.callCount(testFeatures.lsp.sendProgress, mockChatResponseList.length)
863-
assert.deepStrictEqual(chatResult, expectedCompleteChatResult)
857+
assert.deepStrictEqual(chatResult, {
858+
additionalMessages: [],
859+
body: '\n\nHello World!',
860+
messageId: 'mock-message-id',
861+
})
864862
})
865863

866864
it('can use 0 as progress token', async () => {
@@ -872,7 +870,11 @@ describe('AgenticChatController', () => {
872870
const chatResult = await chatResultPromise
873871

874872
sinon.assert.callCount(testFeatures.lsp.sendProgress, mockChatResponseList.length)
875-
assert.deepStrictEqual(chatResult, expectedCompleteChatResult)
873+
assert.deepStrictEqual(chatResult, {
874+
additionalMessages: [],
875+
body: '\n\nHello World!',
876+
messageId: 'mock-message-id',
877+
})
876878
})
877879

878880
it('returns a ResponseError if sendMessage returns an error', async () => {

‎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)