Skip to content

Commit e0b274f

Browse files
justinmk3justinmk3
authored
fix(paidtier): Upgrade success message is unreliable (#1602)
* fix(paidtier): "Upgrade success" message is unreliable - fix spinner placement. - ensure "upgrade successful" notification happens without needing to switch tabs. * fix(paidtier): don't show "Upgrade" to pro-tier Problem: MONTHLY_REQUEST_COUNT is technically possible for pro-tier users. In that case, it would be confusing to show a "Upgrade to Pro" message. Solution: Check for pro-tier before showing the "Upgrade to Pro" message.
1 parent 400c014 commit e0b274f

File tree

10 files changed

+102
-61
lines changed

10 files changed

+102
-61
lines changed

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -922,12 +922,15 @@ export const createMynahUi = (
922922
// Change the sticky banner to show a progress spinner.
923923
const card: typeof freeTierLimitSticky = {
924924
...(isFreeTierLimitUi ? freeTierLimitSticky : upgradePendingSticky),
925-
icon: 'progress',
925+
}
926+
card.header = {
927+
...card.header,
928+
icon: upgradePendingSticky.header?.icon,
929+
iconStatus: upgradePendingSticky.header?.iconStatus,
926930
}
927931
mynahUi.updateStore(tabId, {
928-
// Show a progress ribbon.
929932
promptInputVisible: true,
930-
promptInputStickyCard: isFreeTierLimitUi ? card : null,
933+
promptInputStickyCard: card,
931934
})
932935
} else if (mode === 'paidtier') {
933936
mynahUi.updateStore(tabId, {

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

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@ export const upgradeQButton: ChatItemButton = {
1313
// https://github.com/aws/mynah-ui/blob/main/src/components/icon/icons/q.svg
1414
// https://github.com/aws/mynah-ui/blob/main/src/components/icon/icons/rocket.svg
1515
// icon: MynahIcons.Q,
16-
description: `Upgrade to ${qProName}`,
1716
text: `Subscribe to ${qProName}`,
17+
// description: `Upgrade to ${qProName}`,
1818
status: 'primary',
1919
disabled: false,
2020
}
@@ -65,12 +65,14 @@ export const freeTierLimitSticky: Partial<ChatItem> = {
6565

6666
export const upgradePendingSticky: Partial<ChatItem> = {
6767
messageId: 'upgrade-pending-banner',
68-
body: 'Waiting for subscription status...',
69-
status: 'info',
70-
buttons: [],
68+
body: freeTierLimitSticky.body,
69+
buttons: [upgradeQButton],
70+
header: {
71+
icon: 'progress',
72+
iconStatus: undefined,
73+
body: '### Waiting for subscription status...',
74+
},
7175
canBeDismissed: true,
72-
icon: 'progress',
73-
// iconStatus: 'info',
7476
}
7577

7678
export const upgradeSuccessSticky: Partial<ChatItem> = {
@@ -137,7 +139,7 @@ export const paidTierPromptInput: TextBasedFormItem = {
137139
placeholder: '111111111111',
138140
type: 'textinput',
139141
id: 'paid-tier',
140-
tooltip: `Upgrade to ${qProName}`,
142+
// tooltip: `Upgrade to ${qProName}`,
141143
value: 'true',
142144
icon: 'magic',
143145
}

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

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ import {
8787
getHttpStatusCode,
8888
getRequestID,
8989
getSsoConnectionType,
90-
isFreeTierLimitError,
90+
isUsageLimitError,
9191
isNullish,
9292
} from '../../shared/utils'
9393
import { HELP_MESSAGE, loadingMessage } from '../chat/constants'
@@ -1404,7 +1404,7 @@ export class AgenticChatController implements ChatHandlers {
14041404
*/
14051405
isUserAction(err: unknown, token?: CancellationToken, session?: ChatSessionService): boolean {
14061406
return (
1407-
!isFreeTierLimitError(err) &&
1407+
!isUsageLimitError(err) &&
14081408
(CancellationError.isUserCancelled(err) ||
14091409
err instanceof ToolApprovalException ||
14101410
isRequestAbortedError(err) ||
@@ -2113,12 +2113,14 @@ export class AgenticChatController implements ChatHandlers {
21132113
metric.metric.cwsprChatConversationId = conversationId
21142114
await this.#telemetryController.emitAddMessageMetric(tabId, metric.metric, 'Failed')
21152115

2116-
if (isFreeTierLimitError(err)) {
2117-
this.setPaidTierMode(tabId, 'freetier-limit')
2116+
if (isUsageLimitError(err)) {
2117+
if (this.#paidTierMode !== 'paidtier') {
2118+
this.setPaidTierMode(tabId, 'freetier-limit')
2119+
}
21182120
return new ResponseError<ChatResult>(LSPErrorCodes.RequestFailed, err.message, {
21192121
type: 'answer',
2120-
body: `AmazonQFreeTierLimitError: Free tier limit reached. ${requestID ? `\n\nRequest ID: ${requestID}` : ''}`,
2121-
messageId: 'freetier-limit',
2122+
body: `AmazonQUsageLimitError: Monthly limit reached. ${requestID ? `\n\nRequest ID: ${requestID}` : ''}`,
2123+
messageId: 'monthly-usage-limit',
21222124
buttons: [],
21232125
})
21242126
}
@@ -2160,7 +2162,7 @@ export class AgenticChatController implements ChatHandlers {
21602162
return createAuthFollowUpResult(authFollowType)
21612163
}
21622164

2163-
if (isFreeTierLimitError(err) || customerFacingErrorCodes.includes(err.code)) {
2165+
if (isUsageLimitError(err) || customerFacingErrorCodes.includes(err.code)) {
21642166
this.#features.logging.error(`${loggingUtils.formatErr(err)}`)
21652167
if (err.code === 'InputTooLong') {
21662168
// Clear the chat history in the database for this tab
@@ -2642,18 +2644,20 @@ export class AgenticChatController implements ChatHandlers {
26422644
* 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.
26432645
*
26442646
* `mode` behavior:
2645-
* - 'freetier': treated as 'freetier-limit' if `this.#paidTierMode='freetier-limit'`.
2646-
* - 'freetier-limit': also show "Free Tier limit reached" card in chat.
2647-
* - This mode is "sticky" until 'paidtier' is passed to override it.
2648-
* - 'paidtier': disable any "free-tier limit" UI.
2647+
* - 'freetier': chat-ui clears "limit reached" UI, if any.
2648+
* - 'freetier-limit':
2649+
* - client (IDE) shows a message.
2650+
* - chat-ui shows a chat card.
2651+
* - 'paidtier': disable any "free-tier limit" UI.
2652+
* - 'upgrade-pending': chat-ui shows a progress spinner.
26492653
*/
26502654
setPaidTierMode(tabId?: string, mode?: PaidTierMode) {
26512655
const isBuilderId = getSsoConnectionType(this.#features.credentialsProvider) === 'builderId'
26522656
if (!isBuilderId) {
26532657
return
26542658
}
26552659

2656-
if (this.#paidTierMode === 'freetier-limit' && mode === 'freetier') {
2660+
if (mode === 'freetier' && this.#paidTierMode === 'freetier-limit') {
26572661
// mode = 'freetier-limit' // Sticky while 'freetier'.
26582662
} else if (mode === 'freetier-limit' && mode !== this.#paidTierMode) {
26592663
this.showFreeTierLimitMsgOnClient(tabId)
@@ -2739,7 +2743,7 @@ export class AgenticChatController implements ChatHandlers {
27392743
this.#log('onManageSubscription: missing encodedVerificationUrl in server response')
27402744
this.#features.lsp.window
27412745
.showMessage({
2742-
message: 'Subscription request failed. Check the account id.',
2746+
message: 'Subscription request failed.',
27432747
type: MessageType.Error,
27442748
})
27452749
.catch(e => {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ type AgenticChatErrorCode =
66
| 'FailedResult' // general error when processing tool results
77
| 'InputTooLong' // too much context given to backend service.
88
| 'PromptCharacterLimit' // customer prompt exceeds
9-
| 'AmazonQFreeTierLimitError' // Free Tier limit was reached.
9+
| 'AmazonQUsageLimitError' // Monthly usage limit was reached (usually free-tier user).
1010
| 'ResponseProcessingTimeout' // response didn't finish streaming in the allowed time
1111
| 'MCPServerInitTimeout' // mcp server failed to start within allowed time
1212
| 'MCPToolExecTimeout' // mcp tool call failed to complete within allowed time
@@ -17,7 +17,7 @@ export const customerFacingErrorCodes: AgenticChatErrorCode[] = [
1717
'QModelResponse',
1818
'InputTooLong',
1919
'PromptCharacterLimit',
20-
'AmazonQFreeTierLimitError',
20+
'AmazonQUsageLimitError',
2121
]
2222

2323
export const unactionableErrorCodes: Partial<Record<AgenticChatErrorCode, string>> = {

‎server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,7 @@ import { AgenticChatError, isInputTooLongError, isRequestAbortedError, wrapError
1515
import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager'
1616
import { loggingUtils } from '@aws/lsp-core'
1717
import { Logging } from '@aws/language-server-runtimes/server-interface'
18-
import { getRequestID, isFreeTierLimitError } from '../../shared/utils'
19-
import { AmazonQFreeTierLimitError } from '../../shared/amazonQServiceManager/errors'
18+
import { getRequestID, isUsageLimitError } from '../../shared/utils'
2019

2120
export type ChatSessionServiceConfig = CodeWhispererStreamingClientConfig
2221
type FileChange = { before?: string; after?: string }
@@ -164,10 +163,10 @@ export class ChatSessionService {
164163
}
165164

166165
const requestId = getRequestID(e)
167-
if (isFreeTierLimitError(e)) {
166+
if (isUsageLimitError(e)) {
168167
throw new AgenticChatError(
169168
'Request aborted',
170-
'AmazonQFreeTierLimitError',
169+
'AmazonQUsageLimitError',
171170
e instanceof Error ? e : undefined,
172171
requestId
173172
)

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -86,9 +86,9 @@ export class AmazonQServiceConnectionExpiredError extends AmazonQError {
8686
}
8787
}
8888

89-
export class AmazonQFreeTierLimitError extends AmazonQError {
89+
export class AmazonQUsageLimitError extends AmazonQError {
9090
constructor(cause?: unknown, message: string = 'Free tier limit reached.') {
91-
super(message, 'E_AMAZON_Q_FREE_TIER_LIMIT', cause)
92-
this.name = 'AmazonQFreeTierLimitError'
91+
super(message, 'E_AMAZON_Q_USAGE_LIMIT', cause)
92+
this.name = 'AmazonQUsageLimitError'
9393
}
9494
}

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

Lines changed: 37 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import {
66
Logging,
77
SDKInitializator,
88
CancellationToken,
9+
CancellationTokenSource,
910
} from '@aws/language-server-runtimes/server-interface'
1011
import { waitUntil } from '@aws/lsp-core/out/util/timeoutUtils'
1112
import { AWSError, ConfigurationOptions, CredentialProviderChain, Credentials } from 'aws-sdk'
@@ -158,7 +159,9 @@ export class CodeWhispererServiceIAM extends CodeWhispererServiceBase {
158159
export class CodeWhispererServiceToken extends CodeWhispererServiceBase {
159160
client: CodeWhispererTokenClient
160161
/** Debounce createSubscriptionToken by storing the current, pending promise (if any). */
161-
#createSubscriptionTokenPromise: Promise<CodeWhispererTokenClient.CreateSubscriptionTokenResponse> | undefined
162+
#createSubscriptionTokenPromise?: Promise<CodeWhispererTokenClient.CreateSubscriptionTokenResponse>
163+
/** If user clicks "Upgrade" multiple times, cancel the previous wait-promise. */
164+
#waitUntilSubscriptionCancelSource?: CancellationTokenSource
162165

163166
constructor(
164167
private credentialsProvider: CredentialsProvider,
@@ -376,13 +379,21 @@ export class CodeWhispererServiceToken extends CodeWhispererServiceBase {
376379
async createSubscriptionToken(request: CodeWhispererTokenClient.CreateSubscriptionTokenRequest) {
377380
// Debounce.
378381
if (this.#createSubscriptionTokenPromise) {
379-
// this.logging.debug('createSubscriptionTokenPromise: debounced')
380382
return this.#createSubscriptionTokenPromise
381383
}
382384

383385
this.#createSubscriptionTokenPromise = (async () => {
384386
try {
385-
return this.client.createSubscriptionToken(this.withProfileArn(request)).promise()
387+
const r = await this.client.createSubscriptionToken(this.withProfileArn(request)).promise()
388+
if (!r.encodedVerificationUrl) {
389+
this.logging.error(`setpaidtier
390+
request: ${JSON.stringify(request)}
391+
response: ${JSON.stringify(r as any)}
392+
requestId: ${(r as any).$response?.requestId}
393+
httpStatusCode: ${(r as any).$response?.httpResponse?.statusCode}
394+
headers: ${JSON.stringify((r as any).$response?.httpResponse?.headers)}`)
395+
}
396+
return r
386397
} finally {
387398
this.#createSubscriptionTokenPromise = undefined
388399
}
@@ -441,9 +452,27 @@ export class CodeWhispererServiceToken extends CodeWhispererServiceBase {
441452
* Returns true on success, or false on timeout/cancellation.
442453
*/
443454
async waitUntilSubscriptionActive(cancelToken?: CancellationToken): Promise<boolean> {
455+
// If user clicks "Upgrade" multiple times, cancel any pending waitUntil().
456+
if (this.#waitUntilSubscriptionCancelSource) {
457+
this.#waitUntilSubscriptionCancelSource.cancel()
458+
this.#waitUntilSubscriptionCancelSource.dispose()
459+
}
460+
461+
this.#waitUntilSubscriptionCancelSource = new CancellationTokenSource()
462+
463+
// Combine the external cancelToken (if provided) with our internal one.
464+
const combinedToken = cancelToken
465+
? {
466+
isCancellationRequested: () =>
467+
cancelToken.isCancellationRequested ||
468+
this.#waitUntilSubscriptionCancelSource!.token.isCancellationRequested,
469+
}
470+
: this.#waitUntilSubscriptionCancelSource.token
471+
444472
const r = await waitUntil(
445473
async () => {
446-
if (cancelToken?.isCancellationRequested) {
474+
if (combinedToken.isCancellationRequested) {
475+
this.logging.info('waitUntilSubscriptionActive: cancelled')
447476
return false
448477
}
449478
const s = await this.getSubscriptionStatus(true)
@@ -457,7 +486,10 @@ export class CodeWhispererServiceToken extends CodeWhispererServiceBase {
457486
interval: 2000,
458487
truthy: true,
459488
}
460-
)
489+
).finally(() => {
490+
this.#waitUntilSubscriptionCancelSource?.dispose()
491+
this.#waitUntilSubscriptionCancelSource = undefined
492+
})
461493

462494
return !!r
463495
}

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,11 @@ import {
1313
SendMessageCommandOutput as SendMessageCommandOutputQDeveloperStreaming,
1414
} from '@amzn/amazon-q-developer-streaming-client'
1515
import { CredentialsProvider, SDKInitializator, Logging } from '@aws/language-server-runtimes/server-interface'
16-
import { getBearerTokenFromProvider, isFreeTierLimitError } from './utils'
16+
import { getBearerTokenFromProvider, isUsageLimitError } from './utils'
1717
import { ConfiguredRetryStrategy } from '@aws-sdk/util-retry'
1818
import { CredentialProviderChain, Credentials } from 'aws-sdk'
1919
import { clientTimeoutMs } from '../language-server/agenticChat/constants'
20-
import { AmazonQFreeTierLimitError } from './amazonQServiceManager/errors'
20+
import { AmazonQUsageLimitError } from './amazonQServiceManager/errors'
2121

2222
export type SendMessageCommandInput =
2323
| SendMessageCommandInputCodeWhispererStreaming
@@ -104,8 +104,8 @@ export class StreamingClientServiceToken extends StreamingClientServiceBase {
104104

105105
return response
106106
} catch (e) {
107-
if (isFreeTierLimitError(e)) {
108-
throw new AmazonQFreeTierLimitError(e)
107+
if (isUsageLimitError(e)) {
108+
throw new AmazonQUsageLimitError(e)
109109
}
110110
throw e
111111
} finally {
@@ -132,8 +132,8 @@ export class StreamingClientServiceToken extends StreamingClientServiceBase {
132132
return response
133133
} catch (e) {
134134
// TODO add a test for this
135-
if (isFreeTierLimitError(e)) {
136-
throw new AmazonQFreeTierLimitError(e)
135+
if (isUsageLimitError(e)) {
136+
throw new AmazonQUsageLimitError(e)
137137
}
138138
throw e
139139
} finally {

‎server/aws-lsp-codewhisperer/src/shared/utils.test.ts

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ import {
1515
getSsoConnectionType,
1616
getUnmodifiedAcceptedTokens,
1717
isAwsThrottlingError,
18-
isFreeTierLimitError,
18+
isUsageLimitError,
1919
isQuotaExceededError,
2020
isStringOrNull,
2121
safeGet,
@@ -291,18 +291,18 @@ describe('isAwsThrottlingError', function () {
291291
})
292292
})
293293

294-
describe('isFreeTierLimitError', function () {
294+
describe('isMonthlyLimitError', function () {
295295
it('false for non-throttling errors', function () {
296296
const regularError = new Error('Some error')
297-
assert.strictEqual(isFreeTierLimitError(regularError), false)
297+
assert.strictEqual(isUsageLimitError(regularError), false)
298298

299299
const e = new Error()
300300
;(e as any).name = 'AWSError'
301301
;(e as any).message = 'Not a throttling error'
302302
;(e as any).code = 'SomeOtherError'
303303
;(e as any).time = new Date()
304304

305-
assert.strictEqual(isFreeTierLimitError(e), false)
305+
assert.strictEqual(isUsageLimitError(e), false)
306306
})
307307

308308
it('false for throttling errors without MONTHLY_REQUEST_COUNT reason', function () {
@@ -313,18 +313,18 @@ describe('isFreeTierLimitError', function () {
313313
;(throttlingError as any).time = new Date()
314314
;(throttlingError as any).reason = 'SOME_OTHER_REASON'
315315

316-
assert.strictEqual(isFreeTierLimitError(throttlingError), false)
316+
assert.strictEqual(isUsageLimitError(throttlingError), false)
317317
})
318318

319319
it('true for throttling errors with MONTHLY_REQUEST_COUNT reason', function () {
320-
const freeTierLimitError = new Error()
321-
;(freeTierLimitError as any).name = 'ThrottlingException'
322-
;(freeTierLimitError as any).message = 'Free tier limit reached'
323-
;(freeTierLimitError as any).code = 'ThrottlingException'
324-
;(freeTierLimitError as any).time = new Date()
325-
;(freeTierLimitError as any).reason = ThrottlingExceptionReason.MONTHLY_REQUEST_COUNT
326-
327-
assert.strictEqual(isFreeTierLimitError(freeTierLimitError), true)
320+
const usageLimitError = new Error()
321+
;(usageLimitError as any).name = 'ThrottlingException'
322+
;(usageLimitError as any).message = 'Free tier limit reached'
323+
;(usageLimitError as any).code = 'ThrottlingException'
324+
;(usageLimitError as any).time = new Date()
325+
;(usageLimitError as any).reason = ThrottlingExceptionReason.MONTHLY_REQUEST_COUNT
326+
327+
assert.strictEqual(isUsageLimitError(usageLimitError), true)
328328
})
329329
})
330330

0 commit comments

Comments
 (0)