|
| 1 | +/*! |
| 2 | + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 3 | + * SPDX-License-Identifier: Apache-2.0 |
| 4 | + */ |
| 5 | + |
| 6 | +import * as vscode from 'vscode' |
| 7 | +import { SearchParams } from '../shared/vscode/uriHandler' |
| 8 | +import { ExtContext } from '../shared/extensions' |
| 9 | +import { deeplinkConnect } from '../awsService/sagemaker/commands' |
| 10 | +import { telemetry } from '../shared/telemetry/telemetry' |
| 11 | +/** |
| 12 | + * Registers the SMUS deeplink URI handler at path `/connect/smus`. |
| 13 | + * |
| 14 | + * This handler processes deeplink URLs from the SageMaker Unified Studio console |
| 15 | + * to establish remote connections to SMUS spaces. |
| 16 | + * |
| 17 | + * @param ctx Extension context containing the URI handler |
| 18 | + * @returns Disposable for cleanup |
| 19 | + */ |
| 20 | +export function register(ctx: ExtContext) { |
| 21 | + async function connectHandler(params: ReturnType<typeof parseConnectParams>) { |
| 22 | + await telemetry.smus_deeplinkConnect.run(async (span) => { |
| 23 | + span.record(extractTelemetryMetadata(params)) |
| 24 | + |
| 25 | + // WORKAROUND: The ws_url from the startSession API call contains a query parameter |
| 26 | + // 'cell-number' within itself. When the entire deeplink URL is processed by the URI |
| 27 | + // handler, 'cell-number' is parsed as a standalone query parameter at the top level |
| 28 | + // instead of remaining part of the ws_url. This causes the ws_url to lose the |
| 29 | + // cell-number context it needs. To fix this, we manually re-append the cell-number |
| 30 | + // query parameter back to the ws_url to restore the original intended URL structure. |
| 31 | + await deeplinkConnect( |
| 32 | + ctx, |
| 33 | + params.connection_identifier, |
| 34 | + params.session, |
| 35 | + `${params.ws_url}&cell-number=${params['cell-number']}`, // Re-append cell-number to ws_url |
| 36 | + params.token, |
| 37 | + params.domain, |
| 38 | + params.app_type, |
| 39 | + true // isSMUS=true for SMUS connections |
| 40 | + ) |
| 41 | + }) |
| 42 | + } |
| 43 | + |
| 44 | + return vscode.Disposable.from(ctx.uriHandler.onPath('/connect/smus', connectHandler, parseConnectParams)) |
| 45 | +} |
| 46 | + |
| 47 | +/** |
| 48 | + * Parses and validates SMUS deeplink URI parameters. |
| 49 | + * |
| 50 | + * Required parameters: |
| 51 | + * - connection_identifier: Space ARN identifying the SMUS space |
| 52 | + * - domain: Domain ID for the SMUS space (SM AI side) |
| 53 | + * - user_profile: User profile name |
| 54 | + * - session: SSM session ID |
| 55 | + * - ws_url: WebSocket URL for SSM connection (originally contains cell-number as a query param) |
| 56 | + * - cell-number: extracted from ws_url during URI parsing |
| 57 | + * - token: Authentication token |
| 58 | + * |
| 59 | + * Optional parameters: |
| 60 | + * - app_type: Application type (e.g., JupyterLab, CodeEditor) |
| 61 | + * - smus_domain_id: SMUS domain identifier |
| 62 | + * - smus_domain_account_id: SMUS domain account ID |
| 63 | + * - smus_project_id: SMUS project identifier |
| 64 | + * - smus_domain_region: SMUS domain region |
| 65 | + * |
| 66 | + * Note: The ws_url from startSession API originally includes cell-number as a query parameter. |
| 67 | + * However, when the deeplink URL is processed, the URI handler extracts cell-number as a |
| 68 | + * separate top-level parameter. This is why we need to re-append it in the connectHandler. |
| 69 | + * |
| 70 | + * @param query URI query parameters |
| 71 | + * @returns Parsed parameters object |
| 72 | + * @throws Error if required parameters are missing |
| 73 | + */ |
| 74 | +export function parseConnectParams(query: SearchParams) { |
| 75 | + const requiredParams = query.getFromKeysOrThrow( |
| 76 | + 'connection_identifier', |
| 77 | + 'domain', |
| 78 | + 'user_profile', |
| 79 | + 'session', |
| 80 | + 'ws_url', |
| 81 | + 'cell-number', |
| 82 | + 'token' |
| 83 | + ) |
| 84 | + const optionalParams = query.getFromKeys( |
| 85 | + 'app_type', |
| 86 | + 'smus_domain_id', |
| 87 | + 'smus_domain_account_id', |
| 88 | + 'smus_project_id', |
| 89 | + 'smus_domain_region' |
| 90 | + ) |
| 91 | + |
| 92 | + return { ...requiredParams, ...optionalParams } |
| 93 | +} |
| 94 | + |
| 95 | +/** |
| 96 | + * Extracts telemetry metadata from URI parameters and space ARN. |
| 97 | + * |
| 98 | + * @param params Parsed URI parameters |
| 99 | + * @returns Telemetry metadata object |
| 100 | + */ |
| 101 | +function extractTelemetryMetadata(params: ReturnType<typeof parseConnectParams>) { |
| 102 | + // Extract metadata from space ARN |
| 103 | + // ARN format: arn:aws:sagemaker:region:account-id:space/domain-id/space-name |
| 104 | + const arnParts = params.connection_identifier.split(':') |
| 105 | + const resourceParts = arnParts[5]?.split('/') // Gets "space/domain-id/space-name" |
| 106 | + |
| 107 | + const projectRegion = arnParts[3] // region from ARN |
| 108 | + const projectAccountId = arnParts[4] // account-id from ARN |
| 109 | + const domainIdFromArn = resourceParts?.[1] // domain-id from ARN |
| 110 | + const spaceName = resourceParts?.[2] // space-name from ARN |
| 111 | + |
| 112 | + return { |
| 113 | + smusDomainId: params.smus_domain_id, |
| 114 | + smusDomainAccountId: params.smus_domain_account_id, |
| 115 | + smusProjectId: params.smus_project_id, |
| 116 | + smusDomainRegion: params.smus_domain_region, |
| 117 | + smusProjectRegion: projectRegion, |
| 118 | + smusProjectAccountId: projectAccountId, |
| 119 | + smusSpaceKey: domainIdFromArn && spaceName ? `${domainIdFromArn}/${spaceName}` : undefined, |
| 120 | + } |
| 121 | +} |
0 commit comments