Skip to content

Commit d612404

Browse files
justinmk3justinmk3
committed
getSubscriptionStatus, disable chat input
1 parent 4b1c6ac commit d612404

File tree

8 files changed

+149
-40
lines changed

8 files changed

+149
-40
lines changed

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

Lines changed: 12 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -265,7 +265,7 @@ export const createMynahUi = (
265265
// We check if tabMetadata.openTabKey exists - if it does and is set to true, we skip showing welcome messages
266266
// since this indicates we're loading a previous chat session rather than starting a new one.
267267
if (!tabStore?.tabMetadata || !tabStore.tabMetadata.openTabKey) {
268-
defaultTabConfig.chatItems = tabFactory.getChatItems(true, programmingModeCardActive, false, [])
268+
defaultTabConfig.chatItems = tabFactory.getChatItems(true, programmingModeCardActive, [])
269269
}
270270
mynahUi.updateStore(tabId, defaultTabConfig)
271271
messager.onTabAdd(tabId)
@@ -482,7 +482,7 @@ export const createMynahUi = (
482482
// Update the tab defaults to hide the programmer mode card for new tabs
483483
mynahUi.updateTabDefaults({
484484
store: {
485-
chatItems: tabFactory.getChatItems(true, false, false),
485+
chatItems: tabFactory.getChatItems(true, false),
486486
},
487487
})
488488
}
@@ -503,7 +503,7 @@ export const createMynahUi = (
503503
},
504504
},
505505
defaults: {
506-
store: tabFactory.createTab(false, false),
506+
store: tabFactory.createTab(false),
507507
},
508508
config: {
509509
maxTabs: 10,
@@ -543,7 +543,7 @@ export const createMynahUi = (
543543
// This distinction helps maintain consistent tab behavior between fresh conversations and restored sessions.
544544
const createTabId = (openTab?: boolean) => {
545545
const tabId = mynahUi.updateStore('', {
546-
...tabFactory.createTab(disclaimerCardActive, false),
546+
...tabFactory.createTab(disclaimerCardActive),
547547
tabMetadata: { openTabKey: openTab ? true : false },
548548
})
549549
if (tabId === undefined) {
@@ -826,20 +826,16 @@ export const createMynahUi = (
826826
if (params.data?.placeholderText === 'upgrade-q') {
827827
const tabId = params.tabId !== 'xxx' ? params.tabId : getOrCreateTabId()!
828828
const upgradeQMode: 'paidtier' | 'freetier' | 'freetier-limit' = (params as any).upgradeQMode
829+
// TODO: only do this for 'freetier-limit'
830+
const needUpgrade = upgradeQMode === 'freetier' || upgradeQMode === 'freetier-limit'
829831

830-
// const chatItem: ChatItem = {
831-
// type: ChatItemType.DIRECTIVE,
832-
// contentHorizontalAlignment: 'center',
833-
// fullWidth: true,
834-
// body: `Upgrade Q: ${upgradeQMode}`,
835-
// }
836-
// mynahUi.addChatItem(tabId, chatItem)
837-
upgradeQButton.description = `Upgrade Q: ${upgradeQMode}`
838832
mynahUi.updateStore(tabId, {
839-
promptInputButtons: upgradeQMode === 'paidtier' ? [] : [upgradeQButton],
840-
chatItems: upgradeQMode === 'freetier-limit' ? [paidTierCard] : [],
833+
promptInputButtons: needUpgrade ? [upgradeQButton] : [],
834+
promptInputDisabledState: needUpgrade,
841835
})
842-
mynahUi.addChatItem(tabId, paidTierCard)
836+
if (needUpgrade) {
837+
mynahUi.addChatItem(tabId, paidTierCard)
838+
}
843839
return
844840
}
845841

@@ -1027,12 +1023,7 @@ ${params.message}`,
10271023
const tabId = createTabId(true)
10281024
if (tabId) {
10291025
mynahUi.updateStore(tabId, {
1030-
chatItems: tabFactory.getChatItems(
1031-
messages ? false : true,
1032-
programmingModeCardActive,
1033-
false,
1034-
messages
1035-
),
1026+
chatItems: tabFactory.getChatItems(messages ? false : true, programmingModeCardActive, messages),
10361027
})
10371028
messager.onOpenTab(requestId, { tabId })
10381029
} else {

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,10 @@ export class TabFactory {
3737
this.initialTabId = TabFactory.generateUniqueId()
3838
}
3939

40-
public createTab(disclaimerCardActive: boolean, showUpgradeButton: boolean): MynahUIDataModel {
40+
public createTab(disclaimerCardActive: boolean): MynahUIDataModel {
4141
const tabData: MynahUIDataModel = {
4242
...this.getDefaultTabData(),
4343
...(disclaimerCardActive ? { promptInputStickyCard: disclaimerCard } : {}),
44-
promptInputButtons: showUpgradeButton ? [upgradeQButton] : [],
4544
promptInputOptions: this.agenticMode ? [pairProgrammingPromptInput] : [],
4645
cancelButtonWhenLoading: this.agenticMode, // supported for agentic chat only
4746
}
@@ -51,15 +50,15 @@ export class TabFactory {
5150
public getChatItems(
5251
needWelcomeMessages: boolean,
5352
pairProgrammingCardActive: boolean,
54-
paidTierCardActive: boolean,
53+
// paidTierCardActive: boolean,
5554
chatMessages?: ChatMessage[]
5655
): ChatItem[] {
5756
return [
5857
...(this.bannerMessage ? [this.getBannerMessage() as ChatItem] : []),
5958
...(needWelcomeMessages
6059
? [
6160
...(this.agenticMode && pairProgrammingCardActive ? [programmerModeCard] : []),
62-
...(paidTierCardActive ? [paidTierCard] : []),
61+
// ...(paidTierCardActive ? [paidTierCard] : []),
6362
{
6463
type: ChatItemType.ANSWER,
6564
body: `Hi, I'm Amazon Q. I can answer your software development questions.

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ export const paidTierCard: ChatItem = {
1111
messageId: 'paidTierCardId',
1212
fullWidth: true,
1313
canBeDismissed: false,
14-
body: 'You have reached the free tier limit. Upgraded to Amazon Q Pro.\n\n[Learn More...](https://aws.amazon.com/q/pricing/)',
14+
body: 'You have reached the free tier limit. Upgrade to Amazon Q Pro.\n\n[Learn More...](https://aws.amazon.com/q/pricing/)',
1515
}
1616

1717
export const paidTierPromptInput: ChatItemFormItem = {

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

Lines changed: 66 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,45 @@ export class AgenticChatController implements ChatHandlers {
269269
failureReason: 'invalid awsAccountId',
270270
}
271271
}
272+
273+
// Note: intentionally async.
274+
try {
275+
const r = await AmazonQTokenServiceManager.getInstance()
276+
?.getCodewhispererService()
277+
.createSubscriptionToken({
278+
accountId: awsAccountId,
279+
})
280+
if (!r.encodedVerificationUrl) {
281+
this.#log('missing encodedVerificationUrl in server response')
282+
return {
283+
success: false,
284+
failureReason: 'missing encodedVerificationUrl in server response',
285+
}
286+
}
287+
288+
this.#log(`encodedVerificationUrl ${r.encodedVerificationUrl}`)
289+
290+
try {
291+
URI.parse(r.encodedVerificationUrl)
292+
} catch (e) {
293+
this.#log(`invalid encodedVerificationUrl: '${r.encodedVerificationUrl}': ${(e as Error).message}`)
294+
return {
295+
success: false,
296+
failureReason: 'invalid encodedVerificationUrl',
297+
}
298+
}
299+
300+
this.#features.lsp.window.showDocument({
301+
external: true, // Client is expected to open the URL in a web browser.
302+
uri: r.encodedVerificationUrl,
303+
})
304+
} catch (e) {
305+
return {
306+
success: false,
307+
failureReason: 'createSubscriptionToken failed',
308+
}
309+
}
310+
272311
this.setUpgradeQMode(params.tabId, 'paidtier')
273312
return { success: true }
274313
} else {
@@ -2187,6 +2226,9 @@ export class AgenticChatController implements ChatHandlers {
21872226

21882227
onLinkClick() {}
21892228

2229+
/**
2230+
* After the Chat UI (mynah-ui) is ready.
2231+
*/
21902232
async onReady() {
21912233
await this.restorePreviousChats()
21922234
try {
@@ -2244,8 +2286,6 @@ export class AgenticChatController implements ChatHandlers {
22442286
name: ChatTelemetryEventName.EnterFocusConversation,
22452287
data: {},
22462288
})
2247-
2248-
this.setUpgradeQMode(params.tabId)
22492289
}
22502290

22512291
onTabRemove(params: TabRemoveParams) {
@@ -2391,7 +2431,7 @@ export class AgenticChatController implements ChatHandlers {
23912431
}
23922432

23932433
/**
2394-
* Enables or disables the "Upgrade Q" UI in the chat component.
2434+
* Updates the "Upgrade Q" (subscription tier) state of the UI in the chat component. If `mode` is not given, the user's subscription status is checked by calling the Q service.
23952435
*
23962436
* `mode` behavior:
23972437
* - 'freetier': always show "Upgrade Q" button.
@@ -2403,15 +2443,36 @@ export class AgenticChatController implements ChatHandlers {
24032443
tabId?: string,
24042444
mode?: 'paidtier' | 'freetier' | 'freetier-limit' /*, session: ChatSessionService*/
24052445
) {
2446+
this.#log(
2447+
`xxx setUpgradeQMode: mode=${mode} getCodewhispererService=${!!this.#serviceManager?.getCodewhispererService()}`
2448+
)
24062449
if (mode === 'freetier-limit') {
24072450
this.#freeTierLimit = true // Sticky until 'paidtier' is sent.
24082451
} else if (mode === 'paidtier') {
24092452
this.#freeTierLimit = false
24102453
} else if (this.#freeTierLimit && (!mode || mode === 'freetier')) {
24112454
mode = 'freetier-limit'
24122455
} else if (!mode) {
2413-
const isFreeTierUser = getSsoConnectionType(this.#features.credentialsProvider) === 'builderId'
2414-
mode = isFreeTierUser ? 'freetier' : 'paidtier'
2456+
try {
2457+
// Note: intentionally async.
2458+
AmazonQTokenServiceManager.getInstance()
2459+
?.getCodewhispererService()
2460+
.getSubscriptionStatus()
2461+
.then(o => {
2462+
this.#log(`xxx getSubscriptionStatus: ${o.status} ${o.encodedVerificationUrl}`)
2463+
this.setUpgradeQMode(tabId, o.status === 'ACTIVE' ? 'paidtier' : 'freetier')
2464+
})
2465+
.catch(err => {
2466+
this.#log(`xxx getSubscriptionStatus failed: ${JSON.stringify(err)}`)
2467+
})
2468+
2469+
// const isFreeTierUser = getSsoConnectionType(this.#features.credentialsProvider) === 'builderId'
2470+
// mode = isFreeTierUser ? 'freetier' : 'paidtier'
2471+
} catch {
2472+
this.#log(`xxx yucky`)
2473+
}
2474+
2475+
return
24152476
}
24162477

24172478
const o: ChatUpdateParams = {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/A
1313
import { AmazonQWorkspaceConfig } from '../../shared/amazonQServiceManager/configurationUtils'
1414
import { TabBarController } from './tabBarController'
1515
import { AmazonQServiceInitializationError } from '../../shared/amazonQServiceManager/errors'
16-
import { safeGet, getSsoConnectionType } from '../../shared/utils'
16+
import { safeGet } from '../../shared/utils'
1717

1818
export const QAgenticChatServer =
1919
// prettier-ignore

‎server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@ export const QConfigurationServerToken =
112112
token: CancellationToken
113113
): Promise<QConfigurationResponse | void> => {
114114
const section = params.section
115+
logging.info('zzzzzzzz1')
115116

116117
let customizations: Customizations | CustomizationWithMetadata[] = []
117118
let developerProfiles: AmazonQDeveloperProfile[] = []

‎server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/BaseAmazonQServiceManager.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@ export abstract class BaseAmazonQServiceManager<
108108
abstract handleOnUpdateConfiguration(params: UpdateConfigurationParams, token: CancellationToken): Promise<void>
109109

110110
public async handleDidChangeConfiguration(): Promise<void> {
111+
this.logging.debug('xxxx 7 handleDidChangeConfiguration')
111112
if (this.isConfigChangeInProgress) {
112113
this.logging.debug(CONFIGURATION_CHANGE_IN_PROGRESS_MSG)
113114
return
@@ -129,6 +130,7 @@ export abstract class BaseAmazonQServiceManager<
129130
}
130131

131132
protected updateCachedServiceConfig(): void {
133+
this.logging.debug('xxxx 8 updateCachedServiceConfig')
132134
if (this.cachedCodewhispererService) {
133135
const customizationArn = this.configurationCache.getProperty('customizationArn')
134136
this.logging.debug(`Using customization=${customizationArn}`)

‎server/aws-lsp-codewhisperer/src/shared/codeWhispererService.ts‎

Lines changed: 63 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ export interface GenerateSuggestionsResponse {
4646

4747
import CodeWhispererSigv4Client = require('../client/sigv4/codewhisperersigv4client')
4848
import CodeWhispererTokenClient = require('../client/token/codewhispererbearertokenclient')
49+
import { getBearerTokenFromProvider } from './utils'
4950

5051
// Right now the only difference between the token client and the IAM client for codewhsiperer is the difference in function name
5152
// This abstract class can grow in the future to account for any additional changes across the clients
@@ -149,13 +150,18 @@ export class CodeWhispererServiceIAM extends CodeWhispererServiceBase {
149150
}
150151
}
151152

153+
/**
154+
* Hint: to get an instance of this, see `AmazonQTokenServiceManager.getCodewhispererService()`.
155+
*/
152156
export class CodeWhispererServiceToken extends CodeWhispererServiceBase {
153157
client: CodeWhispererTokenClient
158+
/** Debounce getSubscriptionStatus by storing the current, pending promise (if any). */
159+
#getSubscriptionStatusPromise: ReturnType<typeof this.createSubscriptionToken> | undefined
154160

155161
constructor(
156-
credentialsProvider: CredentialsProvider,
162+
private credentialsProvider: CredentialsProvider,
157163
workspace: Workspace,
158-
logging: Logging,
164+
private logging: Logging,
159165
codeWhispererRegion: string,
160166
codeWhispererEndpoint: string,
161167
sdkInitializator: SDKInitializator
@@ -166,18 +172,29 @@ export class CodeWhispererServiceToken extends CodeWhispererServiceBase {
166172
endpoint: this.codeWhispererEndpoint,
167173
onRequestSetup: [
168174
req => {
175+
logging.error(`xxx req=${req.operation}`)
169176
this.trackRequest(req)
170-
req.on('build', ({ httpRequest }) => {
171-
const creds = credentialsProvider.getCredentials('bearer') as BearerCredentials
172-
if (!creds?.token) {
173-
throw new Error('Authorization failed, bearer token is not set')
177+
req.on('build', async ({ httpRequest }) => {
178+
try {
179+
const creds = credentialsProvider.getCredentials('bearer') as BearerCredentials
180+
logging.error(`xxx req=${req.operation} token=${creds?.token}`)
181+
if (!creds?.token) {
182+
throw new Error('Authorization failed, bearer token is not set')
183+
}
184+
httpRequest.headers['Authorization'] = `Bearer ${creds.token}`
185+
httpRequest.headers['x-amzn-codewhisperer-optout'] =
186+
`${!this.shareCodeWhispererContentWithAWS}`
187+
} catch (err) {
188+
this.completeRequest(req)
189+
// throw err
174190
}
175-
httpRequest.headers['Authorization'] = `Bearer ${creds.token}`
176-
httpRequest.headers['x-amzn-codewhisperer-optout'] = `${!this.shareCodeWhispererContentWithAWS}`
177191
})
178192
req.on('complete', () => {
179193
this.completeRequest(req)
180194
})
195+
req.on('error', () => {
196+
this.completeRequest(req)
197+
})
181198
},
182199
],
183200
}
@@ -349,4 +366,42 @@ export class CodeWhispererServiceToken extends CodeWhispererServiceBase {
349366
async listFeatureEvaluations(request: CodeWhispererTokenClient.ListFeatureEvaluationsRequest) {
350367
return this.client.listFeatureEvaluations(this.withProfileArn(request)).promise()
351368
}
369+
370+
/**
371+
* cool api you have there 🥹
372+
*/
373+
async createSubscriptionToken(request: CodeWhispererTokenClient.CreateSubscriptionTokenRequest) {
374+
return this.client.createSubscriptionToken(this.withProfileArn(request)).promise()
375+
}
376+
377+
/**
378+
* Gets the Subscription status of the given user.
379+
*/
380+
async getSubscriptionStatus(): ReturnType<typeof this.createSubscriptionToken> {
381+
// Debounce.
382+
if (this.#getSubscriptionStatusPromise) {
383+
// this.logging.debug('getSubscriptionStatus: debounced')
384+
return this.#getSubscriptionStatusPromise
385+
}
386+
387+
this.#getSubscriptionStatusPromise = (async () => {
388+
try {
389+
this.logging.debug('xxx getSubscriptionStatus')
390+
// const creds = this.credentialsProvider.getCredentials('bearer') as BearerCredentials
391+
// if (!creds?.token) {
392+
// throw new Error('Authorization failed, bearer token is not set')
393+
// }
394+
395+
const resp = await this.createSubscriptionToken({
396+
accountId: '111111111111', // Special dummy account for checking Subscription status.
397+
// clientToken: creds.token,
398+
})
399+
return resp
400+
} finally {
401+
this.#getSubscriptionStatusPromise = undefined
402+
}
403+
})()
404+
405+
return this.#getSubscriptionStatusPromise
406+
}
352407
}

0 commit comments

Comments
 (0)