fix(Sagemaker) : Improve error handling when the Space SSH host check fails

Author: bhavya2109sharma

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,
@@ -301,6 +323,9 @@ async function handleRunningSpaceWithDisabledAccess(
         return
     }
 
+    // Validate SSH config before showing progress
+    await validateSshConfig()
+
     // Enable remote access and connect
     const client = sageMakerClient || new SagemakerClient(node.regionCode)
 
@@ -359,6 +384,9 @@ async function handleStoppedSpace(
     spaceName: string,
     sageMakerClient?: SagemakerClient
 ) {
+    // Validate SSH config before showing progress
+    await validateSshConfig()
+
     const client = sageMakerClient || new SagemakerClient(node.regionCode)
 
     try {
@@ -403,6 +431,9 @@ async function handleRunningSpaceWithEnabledAccess(
     spaceName: string,
     sageMakerClient?: SagemakerClient
 ) {
+    // Validate SSH config before showing progress
+    await validateSshConfig()
+
     return await vscode.window.withProgress(
         {
             location: vscode.ProgressLocation.Notification,

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,13 +90,8 @@ 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
-    }
+    // Note: SSH config validation is done in commands.ts before calling this function
+    // to ensure users can fix config issues before the progress dialog appears
 
     // set envirionment variables
     const vars = getSmSsmEnv(ssm, path.join(ctx.globalStorageUri.fsPath, 'sagemaker-local-server-info.json'))