feat(sagemaker): Merge SageMaker SSH Kiro integration

.gitignore

@@ -47,6 +47,7 @@ packages/toolkit/package.nls.json
 packages/toolkit/resources
 packages/amazonq/package.nls.json
 packages/amazonq/resources
+packages/sagemaker-ssh-kiro/resources
 
 # Icons
 packages/*/resources/fonts/aws-toolkit-icons.woff

aws-toolkit-vscode.code-workspace

@@ -12,6 +12,9 @@
         {
             "path": "packages/amazonq",
         },
+        {
+            "path": "packages/sagemaker-ssh-kiro",
+        },
     ],
     "settings": {
         "typescript.tsdk": "node_modules/typescript/lib",

packages/core/resources/smus-context-template.md

@@ -0,0 +1,132 @@
+---
+inclusion: always
+---
+
+# SageMaker Unified Studio Space Context
+
+This workspace is running on an Amazon SageMaker Unified Studio Space.
+
+## Environment
+- Operating system: Ubuntu-based SageMaker Distribution
+- User: sagemaker-user
+- Home directory: /home/sagemaker-user
+- AWS credentials are available via the container credentials provider (AWS_CONTAINER_CREDENTIALS_RELATIVE_URI)
+- Do NOT hardcode AWS credentials; use the default credential chain (e.g., boto3.Session())
+
+## Project Info
+- ~/README.md contains project-specific configuration such as connection names and available compute resources.
+- ~/shared/README.md contains shared project data catalog and storage information.
+Refer to these files when you need details about the project's connections, databases, or S3 paths.
+
+## Project Library (`sagemaker_studio`)
+The `sagemaker_studio` package is pre-installed and provides access to project resources.
+
+### Project
+```python
+from sagemaker_studio import Project
+project = Project()
+
+project.id
+project.name
+project.iam_role          # project IAM role ARN
+project.kms_key_arn       # project KMS key ARN (if configured)
+project.mlflow_tracking_server_arn  # MLflow ARN (if configured)
+project.s3.root           # project S3 root path
+```
+
+### Connections
+```python
+project.connections                          # list all connections
+project.connection()                         # default IAM connection
+project.connection("redshift")               # named connection
+conn.name, conn.id, conn.iam_role
+conn.physical_endpoints[0].host              # endpoint host
+conn.data                                    # all connection properties
+conn.secret                                  # credentials (dict or string)
+conn.create_client()                         # boto3 client with connection credentials
+conn.create_client("glue")                   # boto3 client for specific service
+```
+
+### Catalogs, Databases, and Tables
+```python
+catalog = project.connection().catalog()     # default catalog
+catalog = project.connection().catalog("catalog_id")
+catalog.databases                            # list databases
+db = catalog.database("my_db")
+db.tables                                    # list tables
+table = db.table("my_table")
+table.columns                                # list columns (name, type)
+```
+
+### SQL Utilities
+```python
+from sagemaker_studio import sqlutils
+
+# DuckDB (local, no connection needed)
+result = sqlutils.sql("SELECT * FROM my_df WHERE id > 1")
+
+# Athena
+result = sqlutils.sql("SELECT * FROM orders", connection_name="project.athena")
+
+# Redshift
+result = sqlutils.sql("SELECT * FROM products", connection_name="project.redshift")
+
+# Parameterized queries
+result = sqlutils.sql(
+    "SELECT * FROM orders WHERE status = :status",
+    parameters={"status": "completed"},
+    connection_name="project.redshift"
+)
+
+# Get SQLAlchemy engine
+engine = sqlutils.get_engine(connection_name="project.redshift")
+```
+
+### DataFrame Utilities
+```python
+from sagemaker_studio import dataframeutils
+import pandas as pd
+
+# Read from catalog table
+df = pd.read_catalog_table(database="my_db", table="my_table")
+
+# Write to catalog table
+df.to_catalog_table(database="my_db", table="my_table")
+
+# S3 Tables catalog
+df = pd.read_catalog_table(
+    database="my_db", table="my_table",
+    catalog="s3tablescatalog/my_catalog"
+)
+```
+
+### Spark Utilities
+```python
+from sagemaker_studio import sparkutils
+
+# Initialize Spark Connect session
+spark = sparkutils.init()
+spark = sparkutils.init(connection_name="my_spark_connection")
+
+# Get Spark options for JDBC connections
+options = sparkutils.get_spark_options("my_redshift_connection")
+df = spark.read.format("jdbc").options(**options).option("dbtable", "my_table").load()
+```
+
+## Compute Options
+- **Local Python**: Runs directly on the Space instance. Use for single-machine Python, ML, and AI workloads.
+- **Apache Spark (AWS Glue / Amazon EMR)**: Use `%%pyspark`, `%%scalaspark`, or `%%sql` cell magics in notebooks. Default Spark connection is `project.spark.compatibility`.
+- **SQL (Athena)**: Use `%%sql project.athena` for Trino SQL queries via Amazon Athena.
+- **SQL (Redshift)**: Use `%%sql project.redshift` if a Redshift connection is available.
+
+## Code Patterns
+- Use `sagemaker_studio.Project()` for project-aware sessions and resource discovery
+- Reference data using S3 URIs in s3://bucket/prefix format
+- Write Spark DataFrames to the project catalog: `df.write.saveAsTable(f"{database}.table_name", format='parquet', mode='overwrite')`
+- SQL query results are available as DataFrames in subsequent cells via the `_` variable
+- Use `sqlutils.sql()` for programmatic SQL execution against any connection
+- Use `pd.read_catalog_table()` / `df.to_catalog_table()` for pandas catalog I/O
+
+## MCP Server Configuration
+- When configuring MCP servers, pass AWS credentials via environment variable expansion:
+  "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI": "${AWS_CONTAINER_CREDENTIALS_RELATIVE_URI}"

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

@@ -13,13 +13,18 @@ import { getLogger } from '../../shared/logger/logger'
 import { SagemakerSpaceNode, tryRefreshNode } from './explorer/sagemakerSpaceNode'
 import { isRemoteWorkspace } from '../../shared/vscode/env'
 import _ from 'lodash'
-import { prepareDevEnvConnection, tryRemoteConnection } from './model'
+import {
+    prepareDevEnvConnection,
+    startRemoteViaSageMakerSshKiro,
+    tryRemoteConnection,
+    useSageMakerSshKiroExtension,
+} from './model'
+import { ensureSageMakerSshKiroExtension } from './sagemakerSshKiroUtils'
 import { ExtContext } from '../../shared/extensions'
 import { SagemakerClient } from '../../shared/clients/sagemaker'
 import { AccessDeniedException } from '@amzn/sagemaker-client'
 import { ToolkitError, isUserCancelledError } from '../../shared/errors'
 import { showConfirmationMessage } from '../../shared/utilities/messages'
-import { RemoteSessionError } from '../../shared/remoteSession'
 import {
     ConnectFromRemoteWorkspaceMessage,
     InstanceTypeInsufficientMemory,
@@ -113,11 +118,11 @@ export async function deeplinkConnect(
     )
 
     getLogger().info(
-        `sm:deeplinkConnect: 
-        domain: ${domain}, 
-        appType: ${appType}, 
-        workspaceName: ${workspaceName}, 
-        namespace: ${namespace}, 
+        `sm:deeplinkConnect:
+        domain: ${domain},
+        appType: ${appType},
+        workspaceName: ${workspaceName},
+        namespace: ${namespace},
         eksClusterArn: ${eksClusterArn}`
     )
 
@@ -153,13 +158,21 @@ export async function deeplinkConnect(
         )
 
         try {
-            await startVscodeRemote(
-                remoteEnv.SessionProcess,
-                remoteEnv.hostname,
-                '/home/sagemaker-user',
-                remoteEnv.vscPath,
-                'sagemaker-user'
-            )
+            const path = '/home/sagemaker-user'
+            const username = 'sagemaker-user'
+
+            if (useSageMakerSshKiroExtension()) {
+                await ensureSageMakerSshKiroExtension(ctx.extensionContext)
+                await startRemoteViaSageMakerSshKiro(
+                    remoteEnv.SessionProcess,
+                    remoteEnv.hostname,
+                    path,
+                    remoteEnv.vscPath,
+                    username
+                )
+            } else {
+                await startVscodeRemote(remoteEnv.SessionProcess, remoteEnv.hostname, path, remoteEnv.vscPath, username)
+            }
         } catch (remoteErr: any) {
             throw new ToolkitError(
                 `Failed to establish remote connection: ${remoteErr.message}. Check Remote-SSH logs for details.`,
@@ -174,9 +187,9 @@ export async function deeplinkConnect(
             isSMUS
         )
 
-        if (![RemoteSessionError.MissingExtension, RemoteSessionError.ExtensionVersionTooLow].includes(err.code)) {
+        if (!isUserCancelledError(err)) {
             void vscode.window.showErrorMessage(
-                `Remote connection failed: ${err.message || 'Unknown error'}. Check Output > Log (Window) for details.`
+                `Remote connection failed: ${err?.message || 'Unknown error'}. Check Output > Log (Window) for details.`
             )
             throw err
         }

packages/core/src/awsService/sagemaker/detached-server/errorPage.ts

@@ -79,12 +79,12 @@ export const ErrorText = {
                 'Your session has expired. Update the credentials associated with the IAM profile or use a valid IAM profile, then try again.',
         },
         [ExceptionType.INTERNAL_FAILURE]: {
-            Title: 'Failed to connect remotely to VSCode',
-            Text: 'Unable to establish remote connection to VSCode. This could be due to several factors. Please try again by clicking the VSCode button. If the problem persists, please contact your admin.',
+            Title: 'Failed to connect remotely to the Space',
+            Text: 'Unable to establish remote connection to the Space. This could be due to several factors. Please try again by clicking the connect button. If the problem persists, please contact your admin.',
         },
         [ExceptionType.RESOURCE_LIMIT_EXCEEDED]: {
             Title: 'Connection limit reached',
-            Text: 'You have 10 active remote connections to this space. Stop an existing connection to start a new one.',
+            Text: 'You have 10 active remote connections to this space. Stop an existing connection to start a new one. If the problem persists, try restarting the Space.',
         },
         [ExceptionType.THROTTLING]: {
             Title: 'Too many connection attempts',