Merge branch 'master' into dependabot/npm_and_yarn/master/vscode-lsp-02ec881dff

Author: laileni-aws

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

@@ -191,32 +191,60 @@ export async function openRemoteConnect(
         void vscode.window.showErrorMessage(ConnectFromRemoteWorkspaceMessage)
         return
     }
+
+    const spaceName = node.spaceApp.SpaceName!
     await tryRefreshNode(node)
+
+    // for Stopped SM spaces - check instance type before showing progress
     if (node.getStatus() === 'Stopped') {
         //  In case of SMUS, we pass in a SM Client and for SM AI, it creates a new SM Client.
         const client = sageMakerClient ? sageMakerClient : new SagemakerClient(node.regionCode)
 
         try {
-            await client.startSpace(node.spaceApp.SpaceName!, node.spaceApp.DomainId!)
+            await client.startSpace(spaceName, node.spaceApp.DomainId!)
             await tryRefreshNode(node)
             const appType = node.spaceApp.SpaceSettingsSummary?.AppType
             if (!appType) {
                 throw new ToolkitError('AppType is undefined for the selected space. Cannot start remote connection.', {
                     code: 'undefinedAppType',
                 })
             }
-            await client.waitForAppInService(node.spaceApp.DomainId!, node.spaceApp.SpaceName!, appType)
-            await tryRemoteConnection(node, ctx)
+
+            // Only start showing progress after instance type validation
+            return await vscode.window.withProgress(
+                {
+                    location: vscode.ProgressLocation.Notification,
+                    cancellable: false,
+                    title: `Connecting to ${spaceName}`,
+                },
+                async (progress) => {
+                    progress.report({ message: 'Starting the space.' })
+                    await client.waitForAppInService(node.spaceApp.DomainId!, spaceName, appType)
+                    await tryRemoteConnection(node, ctx, progress)
+                }
+            )
         } catch (err: any) {
             // Ignore InstanceTypeError since it means the user decided not to use an instanceType with more memory
-            if (err.code !== InstanceTypeError) {
-                throw new ToolkitError(`Remote connection failed: ${(err as Error).message}`, {
-                    cause: err as Error,
-                    code: err.code,
-                })
+            // just return without showing progress
+            if (err.code === InstanceTypeError) {
+                return
             }
+            throw new ToolkitError(`Remote connection failed: ${(err as Error).message}`, {
+                cause: err as Error,
+                code: err.code,
+            })
         }
     } else if (node.getStatus() === 'Running') {
-        await tryRemoteConnection(node, ctx)
+        // For running spaces, show progress
+        return await vscode.window.withProgress(
+            {
+                location: vscode.ProgressLocation.Notification,
+                cancellable: false,
+                title: `Connecting to ${spaceName}`,
+            },
+            async (progress) => {
+                await tryRemoteConnection(node, ctx, progress)
+            }
+        )
     }
 }

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

@@ -27,12 +27,14 @@ const logger = getLogger('sagemaker')
 
 export async function tryRemoteConnection(
     node: SagemakerSpaceNode | SagemakerUnifiedStudioSpaceNode,
-    ctx: vscode.ExtensionContext
+    ctx: vscode.ExtensionContext,
+    progress: vscode.Progress<{ message?: string; increment?: number }>
 ) {
     const spaceArn = (await node.getSpaceArn()) as string
     const isSMUS = node instanceof SagemakerUnifiedStudioSpaceNode
     const remoteEnv = await prepareDevEnvConnection(spaceArn, ctx, 'sm_lc', isSMUS, node)
     try {
+        progress.report({ message: 'Opening remote session' })
         await startVscodeRemote(
             remoteEnv.SessionProcess,
             remoteEnv.hostname,
@@ -226,7 +228,13 @@ export async function removeKnownHost(hostname: string): Promise<void> {
         throw ToolkitError.chain(err, 'Failed to read known_hosts file')
     }
 
-    const updatedLines = lines.filter((line) => !line.split(' ')[0].split(',').includes(hostname))
+    const updatedLines = lines.filter((line) => {
+        const entryHostname = line.split(' ')[0].split(',')
+        // Hostnames in the known_hosts file seem to be always lowercase, but keeping the case-sensitive check just in
+        // case. Originally we were only doing the case-sensitive check which caused users to get a host
+        // identification error when reconnecting to a Space after it was restarted.
+        return !entryHostname.includes(hostname) && !entryHostname.includes(hostname.toLowerCase())
+    })
 
     if (updatedLines.length !== lines.length) {
         try {

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

@@ -211,6 +211,28 @@ describe('SageMaker Model', () => {
             assertLogsContain(`Removed '${hostname}' from known_hosts`, false, 'debug')
         })
 
+        it('removes case-sensitive hostname when entry in known_hosts is lowercase', async function () {
+            const mixedCaseHostname = 'Test.Host.Com'
+
+            sandbox.stub(fs, 'existsFile').resolves(true)
+
+            const inputContent = `test.host.com ssh-rsa AAAA\nsome.other.com ssh-rsa BBBB`
+            const expectedOutput = `some.other.com ssh-rsa BBBB`
+
+            sandbox.stub(fs, 'readFileText').resolves(inputContent)
+            const writeStub = sandbox.stub(fs, 'writeFile').resolves()
+
+            await removeKnownHost(mixedCaseHostname)
+
+            sinon.assert.calledWith(
+                writeStub,
+                path.join(os.homedir(), '.ssh', 'known_hosts'),
+                sinon.match((value: string) => value.trim() === expectedOutput),
+                { atomic: true }
+            )
+            assertLogsContain(`Removed '${mixedCaseHostname}' from known_hosts`, false, 'debug')
+        })
+
         it('handles hostname in comma-separated list', async function () {
             sandbox.stub(fs, 'existsFile').resolves(true)