feat(sagemakerunifiedstudio): Add Connection Magics Selector feature to Jupyter Notebooks (aws#2169)

## Problem

[feat(sagemakerunifiedstudio): Add Connection Magics Selector feature to
Jupyter
Notebooks](aws/aws-toolkit-vscode-staging@f0a4690):
- SageMaker Unified Studio CodeEditor users should have the ability to
easily select their SMUS connection/compute cell magics in the Jupyter
Notebook experience.
- The experience should be available when a SageMaker Unified Studio
space is connected either remotely or in the portal.

[feat(sagemakerunifiedstudio): Add resource metadata utils for SMUS
spaces](aws/aws-toolkit-vscode-staging@55a43ad):
- SageMaker Unified Studio depends on resource metadata present in the
SageMaker space for various functionality. This includes cases where the
space is either remotely connected or in the SMUS web portal.
- Depending features include the Connection Magics Selector and other
remote SageMaker Unified Studio features.
- SageMaker Unified Studio spaces store this resource metadata under the
`/opt/ml/metadata/resource-metadata.json` path. This resource metadata
file contains SageMaker and DataZone metadata associated with the space.

## Solution

[feat(sagemakerunifiedstudio): Add Connection Magics Selector feature to
Jupyter
Notebooks](aws/aws-toolkit-vscode-staging@f0a4690):
- Adding Connection Magics Selector feature for SageMaker Unified Studio
spaces.
- The Connection Magics Selector allows users to easily select their
SMUS connection/compute cell magics in the Jupyter Notebook experience.
- This is achieved by adding two new status bar items to each Jupyter
Notebook cell that when clicked show quick pick options for their SMUS
connections/computes. Once a user chooses a connection/compute, the cell
magics are updated accordingly.
- When applying the new cell magics, automatic language syntax
highlighting is also applied to the modified cell based on the
connection's associated language.

[feat(sagemakerunifiedstudio): Add resource metadata utils for SMUS
spaces](aws/aws-toolkit-vscode-staging@55a43ad):
- Adding a resource metadata helper module for SMUS spaces.
- Adding `getResourceMetadata()` and other methods that allow for easy
resource metadata retrieval for use by SMUS features.

## Screenshots

Jupyter Notebook view (w/ Connection and Compute selectors in bottom
right of notebook cells):

<img width="3840" height="1942" alt="Screenshot 2025-07-30 at 11-30-55
getting_started ipynb — src — SageMaker Code Editor"
src="https://github.com/user-attachments/assets/19d4f441-127a-4258-8b2e-0c24b7f40f7d"
/>


Connection Selector:

<img width="847" height="241" alt="Screenshot 2025-07-30 at 11 31 59 AM"
src="https://github.com/user-attachments/assets/93aedb12-12c1-4aaf-8629-99d74ac4cadd"
/>

Compute Selector:

<img width="986" height="251" alt="Screenshot 2025-07-30 at 11 32 38 AM"
src="https://github.com/user-attachments/assets/f91e1e46-285b-4db8-ad20-5a23b29a33c4"
/>

## Testing Done

[feat(sagemakerunifiedstudio): Add Connection Magics Selector feature to
Jupyter
Notebooks](aws/aws-toolkit-vscode-staging@f0a4690):
- WIP - will update

[feat(sagemakerunifiedstudio): Add resource metadata utils for SMUS
spaces](aws/aws-toolkit-vscode-staging@55a43ad):
- Added unit tests, new unit tests pass:
```
  resourceMetadataUtils
    getResourceMetadata()
      ✔ should return metadata when file exists and is valid JSON
      ✔ should return undefined when file does not exist
      ✔ should return undefined and log error when file contains invalid JSON
      ✔ should return undefined and log error when readFileText throws
      ✔ should cache metadata and not re-read file on subsequent calls
      ✔ should handle metadata with missing optional fields
      ✔ should handle metadata with empty AdditionalMetadata
      ✔ should handle empty JSON file
      ✔ should set initialized flag to true even when initialization fails
      ✔ should handle very large JSON files
      ✔ should handle JSON with unexpected additional fields
      ✔ should handle JSON with undefined values
      ✔ should handle concurrent calls to getResourceMetadata
    resetResourceMetadata()
      ✔ should reset cached metadata and allow re-initialization
    resourceMetadataFileExists()
      ✔ should return true when file exists
      ✔ should return false when file does not exist
      ✔ should return false and log error when fs.existsFile throws
```

---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

---------

Co-authored-by: Boyu <[email protected]>
Co-authored-by: invictus <[email protected]>

Author: 3 people

packages/core/src/sagemakerunifiedstudio/activation.ts

@@ -7,9 +7,15 @@ import * as vscode from 'vscode'
 import { activate as activateConnectionMagicsSelector } from './connectionMagicsSelector/activation'
 import { activate as activateNotebookScheduling } from './notebookScheduling/activation'
 import { activate as activateExplorer } from './explorer/activation'
+import { isSageMaker } from '../shared/extensionUtilities'
+import { initializeResourceMetadata } from './shared/utils/resourceMetadataUtils'
 
 export async function activate(extensionContext: vscode.ExtensionContext): Promise<void> {
-    await activateConnectionMagicsSelector(extensionContext)
+    // Only run when environment is a SageMaker Unified Studio space
+    if (isSageMaker('SMUS')) {
+        await initializeResourceMetadata()
+        await activateConnectionMagicsSelector(extensionContext)
+    }
 
     await activateNotebookScheduling(extensionContext)
 

packages/core/src/sagemakerunifiedstudio/connectionMagicsSelector/activation.ts

@@ -4,7 +4,77 @@
  */
 
 import * as vscode from 'vscode'
+import { Constants } from './models/constants'
+import {
+    getStatusBarProviders,
+    showConnectionQuickPick,
+    showProjectQuickPick,
+    parseNotebookCells,
+} from './commands/commands'
 
+/**
+ * Activates the SageMaker Unified Studio Connection Magics Selector feature.
+ *
+ * @param extensionContext The extension context
+ */
 export async function activate(extensionContext: vscode.ExtensionContext): Promise<void> {
-    // NOOP
+    extensionContext.subscriptions.push(
+        vscode.commands.registerCommand(Constants.CONNECTION_COMMAND, () => showConnectionQuickPick()),
+        vscode.commands.registerCommand(Constants.PROJECT_COMMAND, () => showProjectQuickPick())
+    )
+
+    if ('NotebookEdit' in vscode) {
+        const { connectionProvider, projectProvider, separatorProvider } = getStatusBarProviders()
+
+        extensionContext.subscriptions.push(
+            vscode.notebooks.registerNotebookCellStatusBarItemProvider('jupyter-notebook', connectionProvider),
+            vscode.notebooks.registerNotebookCellStatusBarItemProvider('jupyter-notebook', projectProvider),
+            vscode.notebooks.registerNotebookCellStatusBarItemProvider('jupyter-notebook', separatorProvider)
+        )
+
+        extensionContext.subscriptions.push(
+            vscode.window.onDidChangeActiveNotebookEditor(async () => {
+                await parseNotebookCells()
+            })
+        )
+
+        extensionContext.subscriptions.push(vscode.workspace.onDidChangeTextDocument(handleTextDocumentChange))
+
+        void parseNotebookCells()
+    }
+}
+
+/**
+ * Handles text document changes to update status bar when cells are manually edited
+ */
+function handleTextDocumentChange(event: vscode.TextDocumentChangeEvent): void {
+    if (event.document.uri.scheme !== 'vscode-notebook-cell') {
+        return
+    }
+
+    const editor = vscode.window.activeNotebookEditor
+    if (!editor) {
+        return
+    }
+
+    let changedCell: vscode.NotebookCell | undefined
+    for (let i = 0; i < editor.notebook.cellCount; i++) {
+        const cell = editor.notebook.cellAt(i)
+        if (cell.document.uri.toString() === event.document.uri.toString()) {
+            changedCell = cell
+            break
+        }
+    }
+
+    if (changedCell && changedCell.kind === vscode.NotebookCellKind.Code) {
+        const { notebookStateManager } = require('./services/notebookStateManager')
+
+        notebookStateManager.parseCellMagic(changedCell)
+
+        setTimeout(() => {
+            const { connectionProvider, projectProvider } = getStatusBarProviders()
+            connectionProvider.refreshCellStatusBar()
+            projectProvider.refreshCellStatusBar()
+        }, 100)
+    }
 }

packages/core/src/sagemakerunifiedstudio/connectionMagicsSelector/client/connectedSpaceDataZoneClient.ts

@@ -0,0 +1,109 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { DataZone, ListConnectionsCommandOutput } from '@aws-sdk/client-datazone'
+import { getLogger } from '../../../shared/logger/logger'
+
+/**
+ * Represents a DataZone connection
+ */
+export interface DataZoneConnection {
+    connectionId: string
+    name: string
+    type: string
+    props?: Record<string, any>
+}
+
+/**
+ * DataZone client for use in a SageMaker Unified Studio connected space
+ * Uses the user's current AWS credentials (project role credentials)
+ */
+export class ConnectedSpaceDataZoneClient {
+    private datazoneClient: DataZone | undefined
+    private readonly logger = getLogger()
+
+    constructor(
+        private readonly region: string,
+        private readonly customEndpoint?: string
+    ) {}
+
+    /**
+     * Gets the DataZone client, initializing it if necessary
+     * Uses default AWS credentials from the environment (project role)
+     * Supports custom endpoints for non-production environments
+     */
+    private getDataZoneClient(): DataZone {
+        if (!this.datazoneClient) {
+            try {
+                const clientConfig: any = {
+                    region: this.region,
+                }
+
+                // Use custom endpoint if provided (for non-prod environments)
+                if (this.customEndpoint) {
+                    clientConfig.endpoint = this.customEndpoint
+                    this.logger.debug(
+                        `ConnectedSpaceDataZoneClient: Using custom DataZone endpoint: ${this.customEndpoint}`
+                    )
+                } else {
+                    this.logger.debug(
+                        `ConnectedSpaceDataZoneClient: Using default AWS DataZone endpoint for region: ${this.region}`
+                    )
+                }
+
+                this.logger.debug('ConnectedSpaceDataZoneClient: Creating DataZone client with default credentials')
+                this.datazoneClient = new DataZone(clientConfig)
+                this.logger.debug('ConnectedSpaceDataZoneClient: Successfully created DataZone client')
+            } catch (err) {
+                this.logger.error('ConnectedSpaceDataZoneClient: Failed to create DataZone client: %s', err as Error)
+                throw err
+            }
+        }
+        return this.datazoneClient
+    }
+
+    /**
+     * Lists the connections in a DataZone domain and project
+     * @param domainId The DataZone domain identifier
+     * @param projectId The DataZone project identifier
+     * @returns List of connections
+     */
+    public async listConnections(domainId: string, projectId: string): Promise<DataZoneConnection[]> {
+        try {
+            this.logger.info(
+                `ConnectedSpaceDataZoneClient: Listing connections for domain ${domainId}, project ${projectId}`
+            )
+
+            const datazoneClient = this.getDataZoneClient()
+
+            const response: ListConnectionsCommandOutput = await datazoneClient.listConnections({
+                domainIdentifier: domainId,
+                projectIdentifier: projectId,
+            })
+
+            if (!response.items || response.items.length === 0) {
+                this.logger.info(
+                    `ConnectedSpaceDataZoneClient: No connections found for domain ${domainId}, project ${projectId}`
+                )
+                return []
+            }
+
+            const connections: DataZoneConnection[] = response.items.map((connection) => ({
+                connectionId: connection.connectionId || '',
+                name: connection.name || '',
+                type: connection.type || '',
+                props: connection.props || {},
+            }))
+
+            this.logger.info(
+                `ConnectedSpaceDataZoneClient: Found ${connections.length} connections for domain ${domainId}, project ${projectId}`
+            )
+            return connections
+        } catch (err) {
+            this.logger.error('ConnectedSpaceDataZoneClient: Failed to list connections: %s', err as Error)
+            throw err
+        }
+    }
+}