Merge branch 'feature/stepfunctions-execution' into feature/stepfunctions-execution

Author: laileni-aws

packages/amazonq/src/app/inline/recommendationService.ts

@@ -35,6 +35,8 @@ export interface GetAllRecommendationsOptions {
 }
 
 export class RecommendationService {
+    private logger = getLogger()
+
     constructor(
         private readonly sessionManager: SessionManager,
         private cursorUpdateRecorder?: ICursorUpdateRecorder
@@ -117,7 +119,7 @@ export class RecommendationService {
             }
 
             // Handle first request
-            getLogger().info('Sending inline completion request: %O', {
+            this.logger.info('Sending inline completion request: %O', {
                 method: inlineCompletionWithReferencesRequestType.method,
                 request: {
                     textDocument: request.textDocument,
@@ -174,7 +176,7 @@ export class RecommendationService {
                 }
             }
 
-            getLogger().info('Received inline completion response from LSP: %O', {
+            this.logger.info('Received inline completion response from LSP: %O', {
                 sessionId: result.sessionId,
                 latency: Date.now() - t0,
                 itemCount: result.items?.length || 0,
@@ -190,12 +192,14 @@ export class RecommendationService {
 
             if (result.items.length > 0 && result.items[0].isInlineEdit === false) {
                 if (isTriggerByDeletion) {
+                    this.logger.info(`Suggestions were discarded; reason: triggerByDeletion`)
                     return []
                 }
                 // Completion will not be rendered if an edit suggestion has been active for longer than 1 second
                 if (EditSuggestionState.isEditSuggestionDisplayingOverOneSecond()) {
                     const session = this.sessionManager.getActiveSession()
                     if (!session) {
+                        this.logger.error(`Suggestions were discarded; reason: undefined conflicting session`)
                         return []
                     }
                     const params: LogInlineCompletionSessionResultsParams = {
@@ -213,14 +217,14 @@ export class RecommendationService {
                     }
                     languageClient.sendNotification('aws/logInlineCompletionSessionResults', params)
                     this.sessionManager.clear()
-                    getLogger().info(
-                        'Completion discarded due to active edit suggestion displayed longer than 1 second'
+                    this.logger.info(
+                        'Suggetions were discarded; reason: active edit suggestion displayed longer than 1 second'
                     )
                     return []
                 } else if (EditSuggestionState.isEditSuggestionActive()) {
                     // discard the current edit suggestion if its display time is less than 1 sec
                     await commands.executeCommand('aws.amazonq.inline.rejectEdit', true)
-                    getLogger().info('Discarding active edit suggestion displaying less than 1 second')
+                    this.logger.info('Discarding active edit suggestion displaying less than 1 second')
                 }
             }
 
@@ -244,24 +248,25 @@ export class RecommendationService {
 
             // TODO: question, is it possible that the first request returns empty suggestion but has non-empty next token?
             if (result.partialResultToken) {
+                let logstr = `found non null next token; `
                 if (!isInlineEdit) {
                     // If the suggestion is COMPLETIONS and there are more results to fetch, handle them in the background
-                    getLogger().info(
-                        'Suggestion type is COMPLETIONS. Start fetching for more items if partialResultToken exists.'
-                    )
+                    logstr += 'Suggestion type is COMPLETIONS. Start pulling more items'
                     this.processRemainingRequests(languageClient, request, result, token).catch((error) => {
                         languageClient.warn(`Error when getting suggestions: ${error}`)
                     })
                 } else {
                     // Skip fetching for more items if the suggesion is EDITS. If it is EDITS suggestion, only fetching for more
                     // suggestions when the user start to accept a suggesion.
                     // Save editsStreakPartialResultToken for the next EDITS suggestion trigger if user accepts.
-                    getLogger().info('Suggestion type is EDITS. Skip fetching for more items.')
+                    logstr += 'Suggestion type is EDITS. Skip pulling more items'
                     this.sessionManager.updateActiveEditsStreakToken(result.partialResultToken)
                 }
+
+                this.logger.info(logstr)
             }
         } catch (error: any) {
-            getLogger().error('Error getting recommendations: %O', error)
+            this.logger.error('Error getting recommendations: %O', error)
             // bearer token expired
             if (error.data && error.data.awsErrorCode === 'E_AMAZON_Q_CONNECTION_EXPIRED') {
                 // ref: https://github.com/aws/aws-toolkit-vscode/blob/amazonq/v1.74.0/packages/core/src/codewhisperer/service/inlineCompletionService.ts#L104

packages/core/src/awsService/sagemaker/commands.ts

@@ -86,7 +86,8 @@ export async function deeplinkConnect(
     session: string,
     wsUrl: string,
     token: string,
-    domain: string
+    domain: string,
+    appType?: string
 ) {
     getLogger().debug(
         `sm:deeplinkConnect: connectionIdentifier: ${connectionIdentifier} session: ${session} wsUrl: ${wsUrl} token: ${token}`
@@ -107,7 +108,8 @@ export async function deeplinkConnect(
             session,
             wsUrl,
             token,
-            domain
+            domain,
+            appType
         )
 
         await startVscodeRemote(

packages/core/src/awsService/sagemaker/credentialMapping.ts

@@ -96,15 +96,16 @@ export async function persistSSMConnection(
     domain: string,
     session?: string,
     wsUrl?: string,
-    token?: string
+    token?: string,
+    appType?: string
 ): Promise<void> {
     const { region } = parseArn(spaceArn)
     const endpoint = DevSettings.instance.get('endpoints', {})['sagemaker'] ?? ''
 
-    // TODO: Hardcoded to 'jupyterlab' due to a bug in Studio that only supports refreshing
-    // the token for both CodeEditor and JupyterLab Apps in the jupyterlab subdomain.
-    // This will be fixed shortly after NYSummit launch to support refresh URL in CodeEditor subdomain.
-    const appSubDomain = 'jupyterlab'
+    let appSubDomain = 'jupyterlab'
+    if (appType && appType.toLowerCase() === 'codeeditor') {
+        appSubDomain = 'code-editor'
+    }
 
     let envSubdomain: string
 

packages/core/src/awsService/sagemaker/model.ts

@@ -56,7 +56,8 @@ export async function prepareDevEnvConnection(
     session?: string,
     wsUrl?: string,
     token?: string,
-    domain?: string
+    domain?: string,
+    appType?: string
 ) {
     const remoteLogger = configureRemoteConnectionLogger()
     const { ssm, vsc, ssh } = (await ensureDependencies()).unwrap()
@@ -82,7 +83,7 @@ export async function prepareDevEnvConnection(
             await persistSmusProjectCreds(spaceArn, node as SagemakerUnifiedStudioSpaceNode)
         }
     } else if (connectionType === 'sm_dl') {
-        await persistSSMConnection(spaceArn, domain ?? '', session, wsUrl, token)
+        await persistSSMConnection(spaceArn, domain ?? '', session, wsUrl, token, appType)
     }
 
     await startLocalServer(ctx)

packages/core/src/awsService/sagemaker/uriHandlers.ts

@@ -18,7 +18,8 @@ export function register(ctx: ExtContext) {
                 params.session,
                 `${params.ws_url}&cell-number=${params['cell-number']}`,
                 params.token,
-                params.domain
+                params.domain,
+                params.app_type
             )
         })
     }
@@ -27,7 +28,7 @@ export function register(ctx: ExtContext) {
 }
 
 export function parseConnectParams(query: SearchParams) {
-    const params = query.getFromKeysOrThrow(
+    const requiredParams = query.getFromKeysOrThrow(
         'connection_identifier',
         'domain',
         'user_profile',
@@ -36,5 +37,7 @@ export function parseConnectParams(query: SearchParams) {
         'cell-number',
         'token'
     )
-    return params
+    const optionalParams = query.getFromKeys('app_type')
+
+    return { ...requiredParams, ...optionalParams }
 }

packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts

@@ -44,6 +44,7 @@ describe('SageMaker URI handler', function () {
             ws_url: 'wss://example.com',
             'cell-number': '4',
             token: 'my-token',
+            app_type: 'jupyterlab',
         }
 
         const uri = createConnectUri(params)
@@ -55,5 +56,24 @@ describe('SageMaker URI handler', function () {
         assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[3], 'wss://example.com&cell-number=4')
         assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[4], 'my-token')
         assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[5], 'my-domain')
+        assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[6], 'jupyterlab')
+    })
+
+    it('calls deeplinkConnect with undefined app_type when not provided', async function () {
+        const params = {
+            connection_identifier: 'abc123',
+            domain: 'my-domain',
+            user_profile: 'me',
+            session: 'sess-xyz',
+            ws_url: 'wss://example.com',
+            'cell-number': '4',
+            token: 'my-token',
+        }
+
+        const uri = createConnectUri(params)
+        await handler.handleUri(uri)
+
+        assert.ok(deeplinkConnectStub.calledOnce)
+        assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[6], undefined)
     })
 })