@@ -262,7 +262,7 @@ export class AgenticChatController implements ChatHandlers {
262262 this . #stoppedToolUses. add ( params . messageId )
263263 await this . #renderStoppedShellCommand( params . tabId , params . messageId )
264264 return { success : true }
265- } else if ( params . buttonId === 'upgrade-q' ) {
265+ } else if ( params . buttonId === 'paidtier- upgrade-q' ) {
266266 const awsAccountId = ( params as any ) . awsAccountId
267267 if ( typeof awsAccountId !== 'string' ) {
268268 this . #log( `invalid awsAccountId: ${ awsAccountId } )
@@ -272,63 +272,14 @@ export class AgenticChatController implements ChatHandlers {
272272 }
273273 }
274274
275- // Note: intentionally async.
276- try {
277- const r = await AmazonQTokenServiceManager . getInstance ( )
278- ?. getCodewhispererService ( )
279- . createSubscriptionToken ( {
280- accountId : awsAccountId ,
281- } )
282-
283- if ( ! r . encodedVerificationUrl ) {
284- this . #log( 'missing encodedVerificationUrl in server response' )
285- this . #features. lsp . window
286- . showMessage ( {
287- message : 'Subscription request failed. Check the account id.' ,
288- type : MessageType . Error ,
289- } )
290- . catch ( e => {
291- this . #log( `showMessage failed: ${ ( e as Error ) . message } )
292- } )
293- return {
294- success : false ,
295- failureReason : 'missing encodedVerificationUrl in server response' ,
296- }
297- }
298-
299- const uri = r . encodedVerificationUrl
300-
301- try {
302- URI . parse ( uri )
303- } catch ( e ) {
304- this . #log( `invalid encodedVerificationUrl: '${ uri } ${ ( e as Error ) . message } )
305- return {
306- success : false ,
307- failureReason : 'invalid encodedVerificationUrl' ,
308- }
309- }
310-
311- this . #features. lsp . window
312- . showMessage ( {
313- message : 'Upgraded to [Amazon Q Pro](https://aws.amazon.com/q/)' ,
314- type : MessageType . Info ,
315- } )
316- . catch ( e => {
317- this . #log( `showMessage failed: ${ ( e as Error ) . message } )
318- } )
319-
320- this . #features. lsp . window . showDocument ( {
321- external : true , // Client is expected to open the URL in a web browser.
322- uri : uri ,
323- } )
324- } catch ( e ) {
275+ const errmsg = await this . onPaidTierUpgradeClicked ( params . tabId , awsAccountId )
276+ if ( errmsg !== '' ) {
325277 return {
326278 success : false ,
327- failureReason : 'createSubscriptionToken failed' ,
279+ failureReason : errmsg ,
328280 }
329281 }
330282
331- this . setPaidTierMode ( params . tabId , 'paidtier-success' )
332283 return { success : true }
333284 } else {
334285 return {
@@ -2112,6 +2063,8 @@ export class AgenticChatController implements ChatHandlers {
21122063 name : ChatTelemetryEventName . EnterFocusConversation ,
21132064 data : { } ,
21142065 } )
2066+
2067+ this . setPaidTierMode ( params . tabId )
21152068 }
21162069
21172070 onTabRemove ( params : TabRemoveParams ) {
@@ -2279,32 +2232,30 @@ export class AgenticChatController implements ChatHandlers {
22792232 * - 'paidtier': disable any "free-tier limit" UI.
22802233 */
22812234 setPaidTierMode ( tabId ?: string , mode ?: PaidTierMode ) {
2282- this . #log( `xxx setPaidTierMode: mode=${ mode } )
2283-
2284- if ( mode === 'freetier-limit' ) {
2285- this . #paidTierMode = mode // Sticky until 'paidtier' is sent.
2286- } else if ( mode === 'paidtier' ) {
2235+ if ( this . #paidTierMode === 'freetier-limit' && mode === 'freetier' ) {
2236+ mode = 'freetier-limit' // Sticky while 'freetier'.
2237+ } else if ( mode ) {
22872238 this . #paidTierMode = mode
2288- } else if ( this . #paidTierMode === 'freetier-limit' && mode === 'freetier' ) {
2289- mode = 'freetier-limit'
2290- } else if ( ! mode ) {
2239+ } else {
22912240 // Note: intentionally async.
22922241 AmazonQTokenServiceManager . getInstance ( )
2293- ? .getCodewhispererService ( )
2242+ . getCodewhispererService ( )
22942243 . getSubscriptionStatus ( )
22952244 . then ( o => {
2296- this . #log( `xxx getSubscriptionStatus: ${ o . status } ${ o . encodedVerificationUrl } )
2245+ this . #log( `setPaidTierMode: getSubscriptionStatus: ${ o . status } ${ o . encodedVerificationUrl } )
22972246 this . setPaidTierMode ( tabId , o . status === 'ACTIVE' ? 'paidtier' : 'freetier' )
22982247 } )
22992248 . catch ( err => {
2300- this . #log( `xxx getSubscriptionStatus failed: ${ JSON . stringify ( err ) } )
2249+ this . #log( `setPaidTierMode: getSubscriptionStatus failed: ${ JSON . stringify ( err ) } )
23012250 } )
23022251 // const isFreeTierUser = getSsoConnectionType(this.#features.credentialsProvider) === 'builderId'
23032252 // mode = isFreeTierUser ? 'freetier' : 'paidtier'
23042253
23052254 return
23062255 }
23072256
2257+ this . #log( `setPaidTierMode: mode=${ mode } )
2258+
23082259 const o : ChatUpdateParams = {
23092260 tabId : tabId ?? '' ,
23102261 // data: { messages: [] },
@@ -2314,6 +2265,104 @@ export class AgenticChatController implements ChatHandlers {
23142265 this . #features. chat . sendChatUpdate ( o )
23152266 }
23162267
2268+ /**
2269+ * Starts the "Upgrade Q" flow for a free-tier user:
2270+ *
2271+ * 0. `awsAccountId` was provided by the IDE extension.
2272+ * 1. Call `createSubscriptionToken(awsAccountId)`.
2273+ * 2. Set the UI to show "Waiting…" progress indicator.
2274+ * 3. Return result, and...
2275+ * 4. ASYNCHRONOUSLY poll subscription status until success.
2276+ * - Update the UI on success/failure.
2277+ *
2278+ * @param awsAccountId AWS account ID to create subscription for
2279+ * @returns Empty string on success, or error message on failure.
2280+ */
2281+ async onPaidTierUpgradeClicked ( tabId : string , awsAccountId : string ) : Promise < string > {
2282+ if ( typeof awsAccountId !== 'string' ) {
2283+ this . #log( `invalid awsAccountId: ${ awsAccountId } )
2284+ return 'invalid awsAccountId'
2285+ }
2286+
2287+ try {
2288+ const client = AmazonQTokenServiceManager . getInstance ( ) . getCodewhispererService ( )
2289+ const r = await client . createSubscriptionToken ( {
2290+ accountId : awsAccountId ,
2291+ } )
2292+
2293+ if ( ! r . encodedVerificationUrl ) {
2294+ this . #log( 'missing encodedVerificationUrl in server response' )
2295+ this . #features. lsp . window
2296+ . showMessage ( {
2297+ message : 'Subscription request failed. Check the account id.' ,
2298+ type : MessageType . Error ,
2299+ } )
2300+ . catch ( e => {
2301+ this . #log( `showMessage failed: ${ ( e as Error ) . message } )
2302+ } )
2303+ return 'missing encodedVerificationUrl in server response'
2304+ }
2305+
2306+ const uri = r . encodedVerificationUrl
2307+
2308+ try {
2309+ URI . parse ( uri )
2310+ } catch ( e ) {
2311+ this . #log( `invalid encodedVerificationUrl: '${ uri } ${ ( e as Error ) . message } )
2312+ return 'invalid encodedVerificationUrl'
2313+ }
2314+
2315+ this . #log( `createSubscriptionToken status: ${ r . status } ${ uri } )
2316+ // Set UI to "progress" mode.
2317+ this . setPaidTierMode ( tabId , 'freetier-upgrade-pending' )
2318+
2319+ // Navigate user to the browser, where they will complete "Upgrade Q" flow.
2320+ this . #features. lsp . window . showDocument ( {
2321+ external : true , // Client is expected to open the URL in a web browser.
2322+ uri : uri ,
2323+ } )
2324+
2325+ // Now asynchronously wait for the user to complete the "Upgrade Q" flow.
2326+ client
2327+ . waitUntilSubscriptionActive ( )
2328+ . then ( r => {
2329+ if ( r !== true ) {
2330+ this . setPaidTierMode ( tabId , 'freetier' )
2331+
2332+ this . #features. lsp . window
2333+ . showMessage ( {
2334+ message : 'Timeout or cancellation while waiting for Amazon Q subscription' ,
2335+ type : MessageType . Error ,
2336+ } )
2337+ . catch ( e => {
2338+ this . #log( `showMessage failed: ${ ( e as Error ) . message } )
2339+ } )
2340+
2341+ return
2342+ }
2343+
2344+ this . setPaidTierMode ( tabId , 'paidtier-success' )
2345+
2346+ this . #features. lsp . window
2347+ . showMessage ( {
2348+ message : 'Upgraded to [Amazon Q Pro](https://aws.amazon.com/q/)' ,
2349+ type : MessageType . Info ,
2350+ } )
2351+ . catch ( e => {
2352+ this . #log( `showMessage failed: ${ ( e as Error ) . message } )
2353+ } )
2354+ } )
2355+ . catch ( e => {
2356+ this . #log( `waitUntilSubscriptionActive failed: ${ ( e as Error ) . message } )
2357+ } )
2358+
2359+ return ''
2360+ } catch ( e ) {
2361+ this . #log( `createSubscriptionToken failed: ${ ( e as Error ) . message } )
2362+ return 'Failed to create subscription token'
2363+ }
2364+ }
2365+
23172366 async #processGenerateAssistantResponseResponseWithTimeout(
23182367 response : GenerateAssistantResponseCommandOutput ,
23192368 metric : Metric < AddMessageEvent > ,
0 commit comments