From e9415c3f0cc0daeefd33ae91f0b1cf92467acada Mon Sep 17 00:00:00 2001 From: Keyvan Zare Rami Date: Mon, 15 Sep 2025 10:49:09 -0700 Subject: [PATCH 1/2] Add domain acccount id telemetry --- .../explorer/activation.ts | 2 +- .../nodes/sageMakerUnifiedStudioRootNode.ts | 27 ++++++++++++++++++- .../src/shared/telemetry/vscodeTelemetry.json | 20 +++++++++++++- 3 files changed, 46 insertions(+), 3 deletions(-) diff --git a/packages/core/src/sagemakerunifiedstudio/explorer/activation.ts b/packages/core/src/sagemakerunifiedstudio/explorer/activation.ts index 65aed68e670..05e74b8f2e7 100644 --- a/packages/core/src/sagemakerunifiedstudio/explorer/activation.ts +++ b/packages/core/src/sagemakerunifiedstudio/explorer/activation.ts @@ -95,7 +95,7 @@ export async function activate(extensionContext: vscode.ExtensionContext): Promi if (!validateNode(node)) { return } - await telemetry.smus_startSpace.run(async (span) => { + await telemetry.smus_openRemoteConnection.run(async (span) => { span.record({ smusSpaceKey: node.resource.DomainSpaceKey, smusDomainRegion: node.resource.regionCode, diff --git a/packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.ts b/packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.ts index a72db66ea69..3bcb0974222 100644 --- a/packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.ts +++ b/packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.ts @@ -8,6 +8,7 @@ import { TreeNode } from '../../../shared/treeview/resourceTreeDataProvider' import { getIcon } from '../../../shared/icons' import { getLogger } from '../../../shared/logger/logger' import { DataZoneClient, DataZoneProject } from '../../shared/client/datazoneClient' +import { DefaultStsClient } from '../../../shared/clients/stsClient' import { Commands } from '../../../shared/vscode/commands2' import { telemetry } from '../../../shared/telemetry/telemetry' import { createQuickPick } from '../../../shared/ui/pickerPrompter' @@ -256,7 +257,7 @@ export const smusLoginCommand = Commands.declare('aws.smus.login', () => async ( }) } - // Extract domain ID and region for logging + // Extract domain account ID, domain ID, and region for logging const domainId = connection.domainId const region = connection.ssoRegion @@ -265,6 +266,18 @@ export const smusLoginCommand = Commands.declare('aws.smus.login', () => async ( smusDomainId: domainId, awsRegion: region, }) + try { + const derCredProvider = await authProvider.getDerCredentialsProvider() + const stsClient = new DefaultStsClient(region, await derCredProvider.getCredentials()) + const callerIdentity = await stsClient.getCallerIdentity() + span.record({ + smusDomainAccountId: callerIdentity.Account, + }) + } catch (err) { + logger.error( + `Failed to resolve AWS account ID via STS Client for domain ${domainId} in region ${region}: ${err}` + ) + } // Show success message void vscode.window.showInformationMessage( @@ -334,6 +347,18 @@ export const smusSignOutCommand = Commands.declare('aws.smus.signOut', () => asy smusDomainId: domainId, awsRegion: region, }) + try { + const derCredProvider = await authProvider.getDerCredentialsProvider() + const stsClient = new DefaultStsClient(region!, await derCredProvider.getCredentials()) + const callerIdentity = await stsClient.getCallerIdentity() + span.record({ + smusDomainAccountId: callerIdentity.Account, + }) + } catch (err) { + logger.error( + `Failed to resolve AWS account ID via STS Client for domain ${domainId} in region ${region}: ${err}` + ) + } // Delete the connection (this will also invalidate tokens and clear cache) if (activeConnection) { diff --git a/packages/core/src/shared/telemetry/vscodeTelemetry.json b/packages/core/src/shared/telemetry/vscodeTelemetry.json index aefb6e3a1b1..741e22cc38c 100644 --- a/packages/core/src/shared/telemetry/vscodeTelemetry.json +++ b/packages/core/src/shared/telemetry/vscodeTelemetry.json @@ -295,6 +295,16 @@ "name": "smusConnectionType", "type": "string", "description": "SMUS connection type" + }, + { + "name": "smusDomainAccountId", + "type": "string", + "description": "SMUS domain account id" + }, + { + "name": "smusProjectAccountId", + "type": "string", + "description": "SMUS project account id" } ], "metrics": [ @@ -1359,6 +1369,10 @@ { "type": "smusDomainId", "required": false + }, + { + "type": "smusDomainAccountId", + "required": false } ] }, @@ -1369,6 +1383,10 @@ { "type": "smusDomainId", "required": false + }, + { + "type": "smusDomainAccountId", + "required": false } ] }, @@ -1414,7 +1432,7 @@ "passive": true }, { - "name": "smus_startSpace", + "name": "smus_openRemoteConnection", "description": "Emitted whenever a user starts a SMUS space", "metadata": [ { From d1efda0757e28c7bf2475e8b4aa2e821c75c5b4f Mon Sep 17 00:00:00 2001 From: Keyvan Zare Rami Date: Mon, 15 Sep 2025 10:49:09 -0700 Subject: [PATCH 2/2] Add domain acccount id telemetry --- .../explorer/nodes/sageMakerUnifiedStudioRootNode.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.ts b/packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.ts index 3bcb0974222..e2039f9fc9e 100644 --- a/packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.ts +++ b/packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.ts @@ -238,7 +238,9 @@ export const smusLoginCommand = Commands.declare('aws.smus.login', () => async ( if (!domainUrl) { // User cancelled logger.debug('User cancelled domain URL input') - return + throw new ToolkitError('User cancelled domain URL input', { + cancelled: true, + }) } // Show a simple status bar message instead of progress dialog