Skip to content

Commit 11070c3

Browse files
vpbhargavvpbhargav
authored
Merge pull request #2303 from vpbhargav/feature/smus-m2
telemetry(smus): add authmode to telemetry
2 parents 085db30 + db167b8 commit 11070c3

File tree

6 files changed

+96
-18
lines changed

6 files changed

+96
-18
lines changed

‎packages/core/src/sagemakerunifiedstudio/auth/providers/smusAuthenticationProvider.ts‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,7 @@ export class SmusAuthenticationProvider {
167167
}
168168
// Update IAM mode context in SMUS space environment
169169
if (getContext('aws.smus.inSmusSpaceEnvironment')) {
170-
void this.initIamModeContextInSpaceEnvironment()
170+
await this.initIamModeContextInSpaceEnvironment()
171171
}
172172

173173
this.onDidChangeEmitter.fire()
@@ -179,22 +179,22 @@ export class SmusAuthenticationProvider {
179179

180180
// Set initial IAM mode context
181181
void (async () => {
182-
const activeConn = this.activeConnection
183-
if (activeConn && 'type' in activeConn && activeConn.type === 'iam') {
184-
const state = this.auth.getStateMemento()
185-
const smusConnections = (state.get('smus.connections') as any) || {}
186-
const connectionMetadata = smusConnections[activeConn.id]
187-
const isIamDomain = connectionMetadata?.isIamDomain || false
188-
await setSmusIamModeContext(isIamDomain)
182+
// Update IAM mode context in SMUS space environment
183+
if (getContext('aws.smus.inSmusSpaceEnvironment')) {
184+
await this.initIamModeContextInSpaceEnvironment()
189185
} else {
190-
await setSmusIamModeContext(false)
186+
const activeConn = this.activeConnection
187+
if (activeConn && 'type' in activeConn && activeConn.type === 'iam') {
188+
const state = this.auth.getStateMemento()
189+
const smusConnections = (state.get('smus.connections') as any) || {}
190+
const connectionMetadata = smusConnections[activeConn.id]
191+
const isIamDomain = connectionMetadata?.isIamDomain || false
192+
await setSmusIamModeContext(isIamDomain)
193+
} else {
194+
await setSmusIamModeContext(false)
195+
}
191196
}
192197
})()
193-
194-
// Update IAM mode context in SMUS space environment
195-
if (getContext('aws.smus.inSmusSpaceEnvironment')) {
196-
void this.initIamModeContextInSpaceEnvironment()
197-
}
198198
}
199199

200200
/**
@@ -243,12 +243,15 @@ export class SmusAuthenticationProvider {
243243
const resourceMetadata = getResourceMetadata()!
244244
if (resourceMetadata.AdditionalMetadata!.DataZoneDomainRegion) {
245245
// Return a mock connection object for SMUS space environment
246-
// This is typed as SmusConnection to satisfy type checking
246+
// Include type property based on IAM mode context for telemetry
247+
// Note: type will be undefined initially until mode is detected
248+
const isIamMode = getContext('aws.smus.isIamMode')
247249
return {
248250
domainId: resourceMetadata.AdditionalMetadata!.DataZoneDomainId!,
249251
ssoRegion: resourceMetadata.AdditionalMetadata!.DataZoneDomainRegion!,
250252
domainUrl: `https://${resourceMetadata.AdditionalMetadata!.DataZoneDomainId!}.sagemaker.${resourceMetadata.AdditionalMetadata!.DataZoneDomainRegion!}.on.aws/`,
251253
id: randomUUID(),
254+
type: isIamMode !== undefined ? (isIamMode ? 'iam' : 'sso') : undefined,
252255
} as any as SmusConnection
253256
} else {
254257
throw new ToolkitError('Domain region not found in metadata file.')

‎packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioProjectNode.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,13 +87,18 @@ export class SageMakerUnifiedStudioProjectNode implements TreeNode {
8787
return telemetry.smus_renderProjectChildrenNode.run(async (span) => {
8888
try {
8989
const isInSmusSpace = getContext('aws.smus.inSmusSpaceEnvironment')
90+
91+
// Get auth mode directly from connection type
92+
const authMode = this.authProvider.activeConnection?.type
93+
9094
const accountId = await this.authProvider.getDomainAccountId()
9195
span.record({
9296
smusToolkitEnv: isInSmusSpace ? 'smus_space' : 'local',
9397
smusDomainId: this.project?.domainId,
9498
smusDomainAccountId: accountId,
9599
smusProjectId: this.project?.id,
96100
smusDomainRegion: this.authProvider.getDomainRegion(),
101+
...(authMode && { smusAuthMode: authMode }),
97102
})
98103

99104
// Skip access check if we're in SMUS space environment (already in project space)

‎packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.ts‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import { isSmusSsoConnection, isSmusIamConnection } from '../../auth/model'
2323
import { getContext } from '../../../shared/vscode/setContext'
2424
import { createDZClientBaseOnDomainMode } from './utils'
2525
import { DataZoneCustomClientHelper } from '../../shared/client/datazoneCustomClientHelper'
26+
import { recordAuthTelemetry } from '../../shared/telemetry'
2627

2728
const contextValueSmusRoot = 'sageMakerUnifiedStudioRoot'
2829
const contextValueSmusLogin = 'sageMakerUnifiedStudioLogin'
@@ -337,6 +338,11 @@ export const smusLoginCommand = Commands.declare('aws.smus.login', (context: vsc
337338
authCompleted = true
338339
}
339340
}
341+
342+
// Record telemetry with connection details after successful login
343+
const domainId = authProvider.getDomainId?.()
344+
const region = authProvider.getDomainRegion?.()
345+
await recordAuthTelemetry(span, authProvider, domainId, region)
340346
} catch (err) {
341347
const isUserCancelled = err instanceof ToolkitError && err.code === SmusErrorCodes.UserCancelled
342348
if (!isUserCancelled) {
@@ -368,9 +374,13 @@ export const smusSignOutCommand = Commands.declare(
368374
return
369375
}
370376

371-
// Get connection details for logging
377+
// Capture connection details BEFORE signing out (for telemetry)
372378
const activeConnection = authProvider.activeConnection
373-
const domainId = authProvider.getDomainId?.() || 'Unknown'
379+
const domainId = authProvider.getDomainId?.()
380+
const region = authProvider.getDomainRegion?.()
381+
382+
// Record telemetry with captured values BEFORE signing out
383+
await recordAuthTelemetry(span, authProvider, domainId, region)
374384

375385
// Sign out from SMUS (behavior depends on connection type)
376386
if (activeConnection) {
@@ -591,6 +601,7 @@ export async function selectSMUSProject(projectNode?: SageMakerUnifiedStudioProj
591601

592602
const accountId = await authProvider.getDomainAccountId()
593603
span.record({
604+
smusAuthMode: authProvider.activeConnection?.type,
594605
smusDomainId: authProvider.getDomainId(),
595606
smusProjectId: (selectedProject as DataZoneProject).id as string | undefined,
596607
smusDomainRegion: authProvider.getDomainRegion(),

‎packages/core/src/sagemakerunifiedstudio/shared/telemetry.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ export async function recordSpaceTelemetry(
5050
}
5151

5252
span.record({
53+
smusAuthMode: authProvider.activeConnection?.type,
5354
smusSpaceKey: node.resource.DomainSpaceKey,
5455
smusDomainRegion: node.resource.regionCode,
5556
smusDomainId: parent?.getAuthProvider()?.getDomainId(),
@@ -75,6 +76,7 @@ export async function recordAuthTelemetry(
7576
const logger = getLogger('smus')
7677

7778
span.record({
79+
smusAuthMode: authProvider.activeConnection?.type,
7880
smusDomainId: domainId,
7981
awsRegion: region,
8082
})
@@ -106,8 +108,11 @@ export async function recordDataConnectionTelemetry(
106108

107109
try {
108110
const isInSmusSpace = getContext('aws.smus.inSmusSpaceEnvironment')
111+
const authProvider = SmusAuthenticationProvider.fromContext()
109112
const accountId = await connectionCredentialsProvider.getDomainAccountId()
113+
110114
span.record({
115+
smusAuthMode: authProvider.activeConnection?.type,
111116
smusToolkitEnv: isInSmusSpace ? 'smus_space' : 'local',
112117
smusDomainId: connection.domainId,
113118
smusDomainAccountId: accountId,

‎packages/core/src/sagemakerunifiedstudio/uriHandlers.ts‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import { SearchParams } from '../shared/vscode/uriHandler'
88
import { ExtContext } from '../shared/extensions'
99
import { deeplinkConnect } from '../awsService/sagemaker/commands'
1010
import { telemetry } from '../shared/telemetry/telemetry'
11+
import { SmusAuthMode } from '../shared/telemetry/telemetry.gen'
1112
/**
1213
* Registers the SMUS deeplink URI handler at path `/connect/smus`.
1314
*
@@ -62,6 +63,7 @@ export function register(ctx: ExtContext) {
6263
* - smus_domain_account_id: SMUS domain account ID
6364
* - smus_project_id: SMUS project identifier
6465
* - smus_domain_region: SMUS domain region
66+
* - smus_auth_mode: Authentication mode (sso or iam)
6567
*
6668
* Note: The ws_url from startSession API originally includes cell-number as a query parameter.
6769
* However, when the deeplink URL is processed, the URI handler extracts cell-number as a
@@ -86,7 +88,8 @@ export function parseConnectParams(query: SearchParams) {
8688
'smus_domain_id',
8789
'smus_domain_account_id',
8890
'smus_project_id',
89-
'smus_domain_region'
91+
'smus_domain_region',
92+
'smus_auth_mode'
9093
)
9194

9295
return { ...requiredParams, ...optionalParams }
@@ -109,6 +112,10 @@ function extractTelemetryMetadata(params: ReturnType<typeof parseConnectParams>)
109112
const domainIdFromArn = resourceParts?.[1] // domain-id from ARN
110113
const spaceName = resourceParts?.[2] // space-name from ARN
111114

115+
// Validate and cast smusAuthMode to the expected type
116+
const authMode = params.smus_auth_mode
117+
const smusAuthMode: SmusAuthMode | undefined = authMode === 'sso' || authMode === 'iam' ? authMode : undefined
118+
112119
return {
113120
smusDomainId: params.smus_domain_id,
114121
smusDomainAccountId: params.smus_domain_account_id,
@@ -117,5 +124,6 @@ function extractTelemetryMetadata(params: ReturnType<typeof parseConnectParams>)
117124
smusProjectRegion: projectRegion,
118125
smusProjectAccountId: projectAccountId,
119126
smusSpaceKey: domainIdFromArn && spaceName ? `${domainIdFromArn}/${spaceName}` : undefined,
127+
smusAuthMode: smusAuthMode,
120128
}
121129
}

‎packages/core/src/shared/telemetry/vscodeTelemetry.json‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,12 @@
319319
"name": "smusProjectAccountId",
320320
"type": "string",
321321
"description": "SMUS project account id"
322+
},
323+
{
324+
"name": "smusAuthMode",
325+
"type": "string",
326+
"allowedValues": ["sso", "iam"],
327+
"description": "SMUS authentication mode (SSO or IAM)"
322328
}
323329
],
324330
"metrics": [
@@ -1387,6 +1393,10 @@
13871393
{
13881394
"type": "smusDomainAccountId",
13891395
"required": false
1396+
},
1397+
{
1398+
"type": "smusAuthMode",
1399+
"required": false
13901400
}
13911401
]
13921402
},
@@ -1401,6 +1411,10 @@
14011411
{
14021412
"type": "smusDomainAccountId",
14031413
"required": false
1414+
},
1415+
{
1416+
"type": "smusAuthMode",
1417+
"required": false
14041418
}
14051419
]
14061420
},
@@ -1423,6 +1437,10 @@
14231437
{
14241438
"type": "smusDomainRegion",
14251439
"required": false
1440+
},
1441+
{
1442+
"type": "smusAuthMode",
1443+
"required": false
14261444
}
14271445
]
14281446
},
@@ -1449,6 +1467,10 @@
14491467
{
14501468
"type": "smusDomainRegion",
14511469
"required": false
1470+
},
1471+
{
1472+
"type": "smusAuthMode",
1473+
"required": false
14521474
}
14531475
],
14541476
"passive": true
@@ -1484,6 +1506,10 @@
14841506
{
14851507
"type": "smusProjectAccountId",
14861508
"required": false
1509+
},
1510+
{
1511+
"type": "smusAuthMode",
1512+
"required": false
14871513
}
14881514
]
14891515
},
@@ -1518,6 +1544,10 @@
15181544
{
15191545
"type": "smusProjectAccountId",
15201546
"required": false
1547+
},
1548+
{
1549+
"type": "smusAuthMode",
1550+
"required": false
15211551
}
15221552
]
15231553
},
@@ -1556,6 +1586,10 @@
15561586
{
15571587
"type": "smusConnectionType",
15581588
"required": false
1589+
},
1590+
{
1591+
"type": "smusAuthMode",
1592+
"required": false
15591593
}
15601594
]
15611595
},
@@ -1594,6 +1628,10 @@
15941628
{
15951629
"type": "smusConnectionType",
15961630
"required": false
1631+
},
1632+
{
1633+
"type": "smusAuthMode",
1634+
"required": false
15971635
}
15981636
]
15991637
},
@@ -1632,6 +1670,10 @@
16321670
{
16331671
"type": "smusConnectionType",
16341672
"required": false
1673+
},
1674+
{
1675+
"type": "smusAuthMode",
1676+
"required": false
16351677
}
16361678
]
16371679
},
@@ -1673,6 +1715,10 @@
16731715
{
16741716
"type": "smusSpaceKey",
16751717
"required": false
1718+
},
1719+
{
1720+
"type": "smusAuthMode",
1721+
"required": false
16761722
}
16771723
]
16781724
}

0 commit comments

Comments
 (0)