Add domain acccount id telemetry

Author: kzr-at-amazon

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,

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) {

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": [
                 {