Skip to content

Commit 9549cde

Browse files
vidkhathvidkhath
committed
feat: implement restore tab
1 parent 0eafb3f commit 9549cde

File tree

8 files changed

+126
-34
lines changed

8 files changed

+126
-34
lines changed

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,7 @@ export const createChat = (
126126
mynahApi.addChatResponse(message.params, message.tabId, message.isPartialResult)
127127
break
128128
case OPEN_TAB_REQUEST_METHOD:
129-
mynahApi.openTab(message.params as OpenTabParams)
129+
mynahApi.openTab(message.requestId, message.params as OpenTabParams)
130130
break
131131
case SEND_TO_PROMPT:
132132
mynahApi.sendToPrompt((message as SendToPromptMessage).params)
@@ -224,9 +224,10 @@ export const createChat = (
224224
disclaimerAcknowledged: () => {
225225
sendMessageToClient({ command: DISCLAIMER_ACKNOWLEDGED })
226226
},
227-
onOpenTab: (params: OpenTabResult | ErrorResult) => {
227+
onOpenTab: (requestId: string, params: OpenTabResult | ErrorResult) => {
228228
if ('tabId' in params) {
229229
sendMessageToClient({
230+
requestId: requestId,
230231
command: OPEN_TAB_REQUEST_METHOD,
231232
params: {
232233
success: true,
@@ -235,6 +236,7 @@ export const createChat = (
235236
})
236237
} else {
237238
sendMessageToClient({
239+
requestId: requestId,
238240
command: OPEN_TAB_REQUEST_METHOD,
239241
params: {
240242
success: false,

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,6 @@ import {
4747
ENTER_FOCUS,
4848
ERROR_MESSAGE_TELEMETRY_EVENT,
4949
EXIT_FOCUS,
50-
FILE_CLICK_TELEMETRY_EVENT,
5150
INFO_LINK_CLICK_TELEMETRY_EVENT,
5251
INSERT_TO_CURSOR_POSITION_TELEMETRY_EVENT,
5352
LINK_CLICK_TELEMETRY_EVENT,
@@ -79,7 +78,7 @@ export interface OutboundChatApi {
7978
infoLinkClick(params: InfoLinkClickParams): void
8079
uiReady(): void
8180
disclaimerAcknowledged(): void
82-
onOpenTab(result: OpenTabResult | ErrorResult): void
81+
onOpenTab(requestId: string, result: OpenTabResult | ErrorResult): void
8382
createPrompt(params: CreatePromptParams): void
8483
fileClick(params: FileClickParams): void
8584
listConversations(params: ListConversationsParams): void
@@ -179,16 +178,15 @@ export class Messager {
179178
this.chatApi.telemetry({ ...params, name: ERROR_MESSAGE_TELEMETRY_EVENT })
180179
}
181180

182-
onOpenTab = (result: OpenTabResult | ErrorResult): void => {
183-
this.chatApi.onOpenTab(result)
181+
onOpenTab = (requestId: string, result: OpenTabResult | ErrorResult): void => {
182+
this.chatApi.onOpenTab(requestId, result)
184183
}
185184

186185
onCreatePrompt = (promptName: string): void => {
187186
this.chatApi.createPrompt({ promptName })
188187
}
189188

190189
onFileClick = (params: FileClickParams): void => {
191-
this.chatApi.telemetry({ ...params, name: FILE_CLICK_TELEMETRY_EVENT })
192190
this.chatApi.fileClick(params)
193191
}
194192

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

Lines changed: 43 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { Messager, OutboundChatApi } from './messager'
66
import { TabFactory } from './tabs/tabFactory'
77
import { ChatItemType, MynahUI } from '@aws/mynah-ui'
88
import { ChatClientAdapter } from '../contracts/chatClientAdapter'
9+
import { ChatMessage } from '@aws/language-server-runtimes-types'
910

1011
describe('MynahUI', () => {
1112
let messager: Messager
@@ -22,6 +23,7 @@ describe('MynahUI', () => {
2223
let onQuickActionSpy: sinon.SinonSpy
2324
let onOpenTabSpy: sinon.SinonSpy
2425
let selectTabSpy: sinon.SinonSpy
26+
const requestId = '1234'
2527

2628
beforeEach(() => {
2729
outboundChatApi = {
@@ -118,12 +120,41 @@ describe('MynahUI', () => {
118120
})
119121

120122
describe('openTab', () => {
121-
it('should create a new tab if tabId not passed', () => {
123+
it('should create a new tab with welcome messages if tabId not passed and previous messages not passed', () => {
122124
createTabStub.resetHistory()
123125

124-
inboundChatApi.openTab({})
126+
inboundChatApi.openTab(requestId, {})
125127

126-
sinon.assert.calledOnceWithExactly(createTabStub, true, false)
128+
sinon.assert.calledOnceWithExactly(createTabStub, true, false, undefined)
129+
sinon.assert.notCalled(selectTabSpy)
130+
sinon.assert.calledOnce(onOpenTabSpy)
131+
})
132+
133+
it('should create a new tab with messages if tabId is not passed and previous messages are passed', () => {
134+
const mockMessages: ChatMessage[] = [
135+
{
136+
messageId: 'msg1',
137+
body: 'Test message 1',
138+
type: ChatItemType.PROMPT,
139+
},
140+
{
141+
messageId: 'msg2',
142+
body: 'Test message 2',
143+
type: ChatItemType.ANSWER,
144+
},
145+
]
146+
147+
createTabStub.resetHistory()
148+
149+
inboundChatApi.openTab(requestId, {
150+
newTabOptions: {
151+
data: {
152+
messages: mockMessages,
153+
},
154+
},
155+
})
156+
157+
sinon.assert.calledOnceWithExactly(createTabStub, false, false, mockMessages)
127158
sinon.assert.notCalled(selectTabSpy)
128159
sinon.assert.calledOnce(onOpenTabSpy)
129160
})
@@ -133,11 +164,11 @@ describe('MynahUI', () => {
133164
updateStoreSpy.restore()
134165
sinon.stub(mynahUi, 'updateStore').returns(undefined)
135166

136-
inboundChatApi.openTab({})
167+
inboundChatApi.openTab(requestId, {})
137168

138-
sinon.assert.calledOnceWithExactly(createTabStub, true, false)
169+
sinon.assert.calledOnceWithExactly(createTabStub, true, false, undefined)
139170
sinon.assert.notCalled(selectTabSpy)
140-
sinon.assert.calledOnceWithMatch(onOpenTabSpy, { type: 'InvalidRequest' })
171+
sinon.assert.calledOnceWithMatch(onOpenTabSpy, requestId, { type: 'InvalidRequest' })
141172
})
142173

143174
it('should open existing tab if tabId passed and tabId not selected', () => {
@@ -146,11 +177,11 @@ describe('MynahUI', () => {
146177
getSelectedTabIdStub.returns('1')
147178
const tabId = '2'
148179

149-
inboundChatApi.openTab({ tabId })
180+
inboundChatApi.openTab(requestId, { tabId })
150181

151182
sinon.assert.notCalled(createTabStub)
152183
sinon.assert.calledOnceWithExactly(selectTabSpy, tabId)
153-
sinon.assert.calledOnceWithExactly(onOpenTabSpy, { tabId })
184+
sinon.assert.calledOnceWithExactly(onOpenTabSpy, requestId, { tabId })
154185
})
155186

156187
it('should not open existing tab if tabId passed but tabId already selected', () => {
@@ -159,11 +190,11 @@ describe('MynahUI', () => {
159190
const tabId = '1'
160191
getSelectedTabIdStub.returns(tabId)
161192

162-
inboundChatApi.openTab({ tabId })
193+
inboundChatApi.openTab(requestId, { tabId })
163194

164195
sinon.assert.notCalled(createTabStub)
165196
sinon.assert.notCalled(selectTabSpy)
166-
sinon.assert.calledOnceWithExactly(onOpenTabSpy, { tabId })
197+
sinon.assert.calledOnceWithExactly(onOpenTabSpy, requestId, { tabId })
167198
})
168199
})
169200

@@ -178,7 +209,7 @@ describe('MynahUI', () => {
178209
getSelectedTabIdStub.returns(undefined)
179210
inboundChatApi.sendGenericCommand({ genericCommand, selection, tabId, triggerType })
180211

181-
sinon.assert.calledOnceWithExactly(createTabStub, false, false)
212+
sinon.assert.calledOnceWithExactly(createTabStub, false, false, undefined)
182213
sinon.assert.calledThrice(updateStoreSpy)
183214
})
184215

@@ -194,7 +225,7 @@ describe('MynahUI', () => {
194225
getSelectedTabIdStub.returns(tabId)
195226
inboundChatApi.sendGenericCommand({ genericCommand, selection, tabId, triggerType })
196227

197-
sinon.assert.calledOnceWithExactly(createTabStub, false, false)
228+
sinon.assert.calledOnceWithExactly(createTabStub, false, false, undefined)
198229
sinon.assert.calledThrice(updateStoreSpy)
199230
})
200231

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

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import {
1313
isValidAuthFollowUpType,
1414
} from '@aws/chat-client-ui-types'
1515
import {
16+
ChatMessage,
1617
ChatResult,
1718
ContextCommand,
1819
ContextCommandParams,
@@ -50,7 +51,7 @@ export interface InboundChatApi {
5051
sendToPrompt(params: SendToPromptParams): void
5152
sendGenericCommand(params: GenericCommandParams): void
5253
showError(params: ErrorParams): void
53-
openTab(params: OpenTabParams): void
54+
openTab(requestId: string, params: OpenTabParams): void
5455
sendContextCommands(params: ContextCommandParams): void
5556
listConversations(params: ListConversationsResult): void
5657
conversationClicked(params: ConversationClickResult): void
@@ -383,8 +384,11 @@ export const createMynahUi = (
383384
return tabId ? mynahUi.getAllTabs()[tabId]?.store : undefined
384385
}
385386

386-
const createTabId = (needWelcomeMessages: boolean = false) => {
387-
const tabId = mynahUi.updateStore('', tabFactory.createTab(needWelcomeMessages, disclaimerCardActive))
387+
const createTabId = (needWelcomeMessages: boolean = false, chatMessages?: ChatMessage[]) => {
388+
const tabId = mynahUi.updateStore(
389+
'',
390+
tabFactory.createTab(needWelcomeMessages, disclaimerCardActive, chatMessages)
391+
)
388392
if (tabId === undefined) {
389393
mynahUi.notify({
390394
content: uiComponentsTexts.noMoreTabsTooltip,
@@ -540,18 +544,19 @@ ${params.message}`,
540544
messager.onError(params)
541545
}
542546

543-
const openTab = ({ tabId }: OpenTabParams) => {
544-
if (tabId) {
545-
if (tabId !== mynahUi.getSelectedTabId()) {
546-
mynahUi.selectTab(tabId)
547+
const openTab = (requestId: string, params: OpenTabParams) => {
548+
if (params.tabId) {
549+
if (params.tabId !== mynahUi.getSelectedTabId()) {
550+
mynahUi.selectTab(params.tabId)
547551
}
548-
messager.onOpenTab({ tabId })
552+
messager.onOpenTab(requestId, { tabId: params.tabId })
549553
} else {
550-
const tabId = createTabId(true)
554+
const messages = params.newTabOptions?.data?.messages
555+
const tabId = createTabId(messages ? false : true, messages)
551556
if (tabId) {
552-
messager.onOpenTab({ tabId })
557+
messager.onOpenTab(requestId, { tabId })
553558
} else {
554-
messager.onOpenTab({
559+
messager.onOpenTab(requestId, {
555560
type: 'InvalidRequest',
556561
message: 'No more tabs available',
557562
})

‎chat-client/src/client/tabs/tabFactory.ts

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
1-
import { ChatItemType, MynahUIDataModel, QuickActionCommandGroup } from '@aws/mynah-ui'
1+
import { ChatItem, ChatItemType, MynahUIDataModel, QuickActionCommandGroup } from '@aws/mynah-ui'
22
import { disclaimerCard } from '../texts/disclaimer'
3+
import { ChatMessage } from '@aws/language-server-runtimes-types'
34

45
export type DefaultTabData = MynahUIDataModel
56

@@ -16,7 +17,11 @@ export class TabFactory {
1617
private quickActionCommands?: QuickActionCommandGroup[]
1718
) {}
1819

19-
public createTab(needWelcomeMessages: boolean, disclaimerCardActive: boolean): MynahUIDataModel {
20+
public createTab(
21+
needWelcomeMessages: boolean,
22+
disclaimerCardActive: boolean,
23+
chatMessages?: ChatMessage[]
24+
): MynahUIDataModel {
2025
const tabData: MynahUIDataModel = {
2126
...this.getDefaultTabData(),
2227
chatItems: needWelcomeMessages
@@ -32,7 +37,9 @@ export class TabFactory {
3237
followUp: this.getWelcomeBlock(),
3338
},
3439
]
35-
: [],
40+
: chatMessages
41+
? (chatMessages as ChatItem[])
42+
: [],
3643
...(disclaimerCardActive ? { promptInputStickyCard: disclaimerCard } : {}),
3744
}
3845
return tabData

‎chat-client/src/contracts/serverContracts.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ export type ServerMessageCommand =
5454

5555
export interface Message {
5656
command: ServerMessageCommand
57+
requestId?: string
5758
}
5859

5960
export interface ServerMessage extends Message {

‎chat-client/src/contracts/telemetry.ts

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ export const LINK_CLICK_TELEMETRY_EVENT = 'linkClick'
1212
export const INFO_LINK_CLICK_TELEMETRY_EVENT = 'infoLinkClick'
1313
export const SOURCE_LINK_CLICK_TELEMETRY_EVENT = 'sourceLinkClick'
1414
export const AUTH_FOLLOW_UP_CLICKED_TELEMETRY_EVENT = 'authFollowupClicked'
15-
export const FILE_CLICK_TELEMETRY_EVENT = 'fileClick'
1615

1716
export enum RelevancyVoteType {
1817
UP = 'upvote',

‎client/vscode/src/chatActivation.ts

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import {
44
AUTH_FOLLOW_UP_CLICKED,
55
CHAT_OPTIONS,
66
COPY_TO_CLIPBOARD,
7+
UiMessageResultParams,
78
} from '@aws/chat-client-ui-types'
89
import {
910
ChatResult,
@@ -20,6 +21,9 @@ import {
2021
contextCommandsNotificationType,
2122
listConversationsRequestType,
2223
conversationClickRequestType,
24+
ErrorCodes,
25+
openTabRequestType,
26+
ResponseError,
2327
} from '@aws/language-server-runtimes/protocol'
2428
import { v4 as uuidv4 } from 'uuid'
2529
import { Uri, Webview, WebviewView, commands, window } from 'vscode'
@@ -165,6 +169,51 @@ export function registerChat(languageClient: LanguageClient, extensionUri: Uri,
165169
})
166170
})
167171

172+
languageClient.onRequest(openTabRequestType.method, async (params, _) => {
173+
const mapErrorType = (type: string | undefined): number => {
174+
switch (type) {
175+
case 'InvalidRequest':
176+
return ErrorCodes.InvalidRequest
177+
case 'InternalError':
178+
return ErrorCodes.InternalError
179+
case 'UnknownError':
180+
default:
181+
return ErrorCodes.UnknownErrorCode
182+
}
183+
}
184+
const requestId = uuidv4()
185+
186+
webviewView.webview.postMessage({
187+
requestId: requestId,
188+
command: openTabRequestType.method,
189+
params: params,
190+
})
191+
const responsePromise = new Promise<UiMessageResultParams | undefined>((resolve, reject) => {
192+
const timeout = setTimeout(() => {
193+
reject(new Error('Request timed out'))
194+
}, 30000)
195+
196+
const disposable = webviewView.webview.onDidReceiveMessage((message: any) => {
197+
if (message.requestId === requestId) {
198+
clearTimeout(timeout)
199+
disposable.dispose()
200+
resolve(message.params)
201+
}
202+
})
203+
})
204+
205+
const result = await responsePromise
206+
207+
if (result?.success) {
208+
return { tabId: result.result.tabId }
209+
} else {
210+
return new ResponseError(
211+
mapErrorType(result?.error.type),
212+
result?.error.message ?? 'No response from client'
213+
)
214+
}
215+
})
216+
168217
webviewView.webview.html = getWebviewContent(webviewView.webview, extensionUri)
169218

170219
registerGenericCommand('aws.sample-vscode-ext-amazonq.explainCode', 'Explain', webviewView.webview)

0 commit comments

Comments
 (0)