fix(sagemaker): improve error handling when the Space SSH host check fails

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

@@ -30,9 +30,28 @@ import {
     SpaceStatus,
 } from './constants'
 import { SagemakerUnifiedStudioSpaceNode } from '../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode'
+import { SageMakerSshConfig } from './sshConfig'
+import { findSshPath } from '../../shared/utilities/pathFind'
 
 const localize = nls.loadMessageBundle()
 
+/**
+ * Validates SSH configuration before starting connection.
+ */
+async function validateSshConfig(): Promise<void> {
+    const sshPath = await findSshPath()
+    if (!sshPath) {
+        throw new ToolkitError(
+            'SSH is required to connect to SageMaker spaces, but was not found.Install SSH to connect to spaces.'
+        )
+    }
+    const sshConfig = new SageMakerSshConfig(sshPath, 'sm_', 'sagemaker_connect')
+    const result = await sshConfig.ensureValid()
+    if (result.isErr()) {
+        throw result.err()
+    }
+}
+
 export async function filterSpaceAppsByDomainUserProfiles(parentNode: SagemakerParentNode): Promise<void> {
     if (parentNode.domainUserProfiles.size === 0) {
         // if parentNode has not been expanded, domainUserProfiles will be empty
@@ -107,6 +126,9 @@ export async function deeplinkConnect(
         return
     }
 
+    // Validate SSH config before attempting connection
+    await validateSshConfig()
+
     try {
         const remoteEnv = await prepareDevEnvConnection(
             connectionIdentifier,
@@ -201,6 +223,9 @@ export async function openRemoteConnect(
         return
     }
 
+    // Validate SSH config before attempting connection
+    await validateSshConfig()
+
     const spaceName = node.spaceApp.SpaceName!
     await tryRefreshNode(node)
 

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

@@ -45,3 +45,19 @@ export const InstanceTypeNotSelectedMessage = (spaceName: string) => {
 
 export const RemoteAccessRequiredMessage =
     'This space requires remote access to be enabled.\nWould you like to restart the space and connect?\nAny unsaved work will be lost.'
+
+// SSH Configuration Error Messages
+export const SshConfigUpdateDeclinedMessage = (configHostName: string, configPath: string) =>
+    `SSH configuration has an outdated ${configHostName} section. Fix your ${configPath} file manually to enable remote connections.`
+
+export const SshConfigOpenedForEditMessage = () =>
+    `SSH configuration file opened for editing. Fix the issue and try connecting again.`
+
+export const SshConfigSyntaxErrorMessage = (configPath: string) =>
+    `SSH configuration has syntax errors in your ${configPath} file. Fix the configuration manually to enable remote connection.`
+
+export const SshConfigRemovalFailedMessage = (configHostName: string) =>
+    `Failed to remove SSH config section for ${configHostName}`
+
+export const SshConfigUpdateFailedMessage = (configPath: string, configHostName: string) =>
+    `Failed to update SSH config section. Fix your ${configPath} file manually or remove the outdated ${configHostName} section.`

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

@@ -8,7 +8,6 @@
 import * as vscode from 'vscode'
 import { sshAgentSocketVariable, startSshAgent, startVscodeRemote } from '../../shared/extensions/ssh'
 import { createBoundProcess, ensureDependencies } from '../../shared/remoteSession'
-import { SshConfig } from '../../shared/sshConfig'
 import * as path from 'path'
 import { persistLocalCredentials, persistSmusProjectCreds, persistSSMConnection } from './credentialMapping'
 import * as os from 'os'
@@ -91,14 +90,6 @@ export async function prepareDevEnvConnection(
     await startLocalServer(ctx)
     await removeKnownHost(hostname)
 
-    const sshConfig = new SshConfig(ssh, 'sm_', 'sagemaker_connect')
-    const config = await sshConfig.ensureValid()
-    if (config.isErr()) {
-        const err = config.err()
-        logger.error(`sagemaker: failed to add ssh config section: ${err.message}`)
-        throw err
-    }
-
     // set envirionment variables
     const vars = getSmSsmEnv(ssm, path.join(ctx.globalStorageUri.fsPath, 'sagemaker-local-server-info.json'))
     logger.info(`connect script logs at ${vars.LOG_FILE_LOCATION}`)