Skip to content

Commit 77bf844

Browse files
nkomonen-amazonnkomonen-amazon
authored
auth: add more info when connecting to auth (#4625)
* telemetry: add source parameter to aws.codeWhisperer.connect External extensions were calling this API to setup a connection but we did not have insight in to who was calling it. Solution, add a `source` parameter so that the caller can identify themselves. Signed-off-by: Nikolas Komonen <[email protected]> * auth: add credentialStartUrl to metric aws_loginWithBrowser Signed-off-by: Nikolas Komonen <[email protected]> * update telemetry Signed-off-by: Nikolas Komonen <[email protected]> --------- Signed-off-by: Nikolas Komonen <[email protected]>
1 parent 17d63d9 commit 77bf844

File tree

7 files changed

+98
-55
lines changed

7 files changed

+98
-55
lines changed

‎docs/api.md‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,13 @@ Details about any publicly accessible functionalities exposed through [extension
88

99
#### `aws.codeWhisperer.connect`
1010

11-
**Signature**: _async (startUrl?: string, region?: string, customizationArn?: string, customizationNamePrefix?: string) => Promise<void>_
11+
**Signature**: _async (source: string, startUrl?: string, region?: string, customizationArn?: string, customizationNamePrefix?: string) => Promise<void>_
1212

1313
Shortcut command to directly connect to Identity Center or prompt start URL entry, as well as set a customization for CodeWhisperer requests.
1414

1515
This command supports the following arguments:
1616

17+
- source: An identifier of the caller of this command. This can be used for something like telemetry.
1718
- startUrl and region. If both arguments are provided they will be used, otherwise the command prompts for them interactively.
1819
- customizationArn: select customization by ARN. If provided, `customizationNamePrefix` is ignored.
1920
- customizationNamePrefix: select customization by prefix, if `customizationArn` is `undefined`.

‎package-lock.json‎

Lines changed: 3 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/core/src/auth/auth.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -172,10 +172,10 @@ export class Auth implements AuthService, ConnectionManager {
172172
if (profile.type === 'sso') {
173173
const provider = this.getSsoTokenProvider(id, profile)
174174

175-
// We create our own span with run() since we want to
176-
// set the isReAuth attribute
175+
// We cannot easily set isReAuth inside the createToken() call,
176+
// so we need to set it here.
177177
await telemetry.aws_loginWithBrowser.run(async span => {
178-
span.record({ isReAuth: true })
178+
span.record({ isReAuth: true, credentialStartUrl: profile.startUrl })
179179
await this.authenticate(id, () => provider.createToken())
180180
})
181181

‎packages/core/src/auth/sso/ssoAccessTokenProvider.ts‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -199,16 +199,25 @@ export class SsoAccessTokenProvider {
199199
return await pollForTokenWithProgress(() => this.oidc.createToken(tokenRequest), authorization)
200200
}
201201

202-
let token: ReturnType<typeof openBrowserAndWaitUntilComplete>
203-
if (telemetry.spans.map(s => s.name).includes('aws_loginWithBrowser')) {
202+
const token = this.withBrowserLoginTelemetry(() => openBrowserAndWaitUntilComplete())
203+
204+
return this.formatToken(await token, registration)
205+
}
206+
207+
/**
208+
* Wraps the given function with telemetry related to the browser login.
209+
*/
210+
private withBrowserLoginTelemetry<T extends (...args: any[]) => any>(func: T): ReturnType<T> {
211+
if (telemetry.spans.some(s => s.name === 'aws_loginWithBrowser')) {
204212
// During certain flows, eg reauthentication, we are already running within a span (run())
205213
// so we don't need to create a new one.
206-
token = openBrowserAndWaitUntilComplete()
207-
} else {
208-
token = telemetry.aws_loginWithBrowser.run(async () => openBrowserAndWaitUntilComplete())
214+
return func()
209215
}
210216

211-
return this.formatToken(await token, registration)
217+
return telemetry.aws_loginWithBrowser.run(span => {
218+
span.record({ credentialStartUrl: this.profile.startUrl })
219+
return func()
220+
})
212221
}
213222

214223
/**

‎packages/core/src/codewhisperer/commands/basicCommands.ts‎

Lines changed: 48 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -144,48 +144,60 @@ export const showSsoSignIn = Commands.declare('aws.codeWhisperer.sso', () => asy
144144
// Shortcut command to directly connect to Identity Center or prompt start URL entry
145145
// It can optionally set a customization too based on given values to match on
146146
export const connectWithCustomization = Commands.declare(
147-
'aws.codeWhisperer.connect',
148-
() => async (startUrl?: string, region?: string, customizationArn?: string, customizationNamePrefix?: string) => {
149-
// This command supports the following arguments:
150-
// * startUrl and region. If both arguments are provided they will be used, otherwise
151-
// the command prompts for them interactively.
152-
// * customizationArn: select customization by ARN. If provided, `customizationNamePrefix` is ignored.
153-
// * customizationNamePrefix: select customization by prefix, if `customizationArn` is `undefined`.
154-
if (startUrl && region) {
155-
await connectToEnterpriseSso(startUrl, region)
156-
} else {
157-
await getStartUrl()
158-
}
147+
{ id: 'aws.codeWhisperer.connect', compositeKey: { 0: 'source' } },
148+
/**
149+
* This command supports the following arguments:
150+
* @param source - an identifier for who used this command. This value is not explicitly used in the function, but is used elsewhere.
151+
* startUrl and region. If both arguments are provided they will be used, otherwise
152+
* the command prompts for them interactively.
153+
* customizationArn: select customization by ARN. If provided, `customizationNamePrefix` is ignored.
154+
* customizationNamePrefix: select customization by prefix, if `customizationArn` is `undefined`.
155+
*/
156+
() =>
157+
async (
158+
source: string,
159+
startUrl?: string,
160+
region?: string,
161+
customizationArn?: string,
162+
customizationNamePrefix?: string
163+
) => {
164+
if (startUrl && region) {
165+
await connectToEnterpriseSso(startUrl, region)
166+
} else {
167+
await getStartUrl()
168+
}
159169

160-
// No customization match information given, exit early.
161-
if (!customizationArn && !customizationNamePrefix) {
162-
return
163-
}
170+
// No customization match information given, exit early.
171+
if (!customizationArn && !customizationNamePrefix) {
172+
return
173+
}
164174

165-
let persistedCustomizations = getPersistedCustomizations()
175+
let persistedCustomizations = getPersistedCustomizations()
166176

167-
// Check if any customizations have already been persisted.
168-
// If not, call `notifyNewCustomizations` to handle it then recheck.
169-
if (persistedCustomizations.length === 0) {
170-
await notifyNewCustomizations()
171-
persistedCustomizations = getPersistedCustomizations()
172-
}
177+
// Check if any customizations have already been persisted.
178+
// If not, call `notifyNewCustomizations` to handle it then recheck.
179+
if (persistedCustomizations.length === 0) {
180+
await notifyNewCustomizations()
181+
persistedCustomizations = getPersistedCustomizations()
182+
}
173183

174-
// If given an ARN, assume a specific customization is desired and find an entry that matches it. Ignores the prefix logic.
175-
// Otherwise if only a prefix is given, find an entry that matches it.
176-
// Backwards compatible with previous implementation.
177-
const match = customizationArn
178-
? persistedCustomizations.find(c => c.arn === customizationArn)
179-
: persistedCustomizations.find(c => c.name?.startsWith(customizationNamePrefix as string))
184+
// If given an ARN, assume a specific customization is desired and find an entry that matches it. Ignores the prefix logic.
185+
// Otherwise if only a prefix is given, find an entry that matches it.
186+
// Backwards compatible with previous implementation.
187+
const match = customizationArn
188+
? persistedCustomizations.find(c => c.arn === customizationArn)
189+
: persistedCustomizations.find(c => c.name?.startsWith(customizationNamePrefix as string))
180190

181-
// If no match is found, nothing to do :)
182-
if (!match) {
183-
getLogger().error(`No customization match found: arn=${customizationArn} prefix=${customizationNamePrefix}`)
184-
return
191+
// If no match is found, nothing to do :)
192+
if (!match) {
193+
getLogger().error(
194+
`No customization match found: arn=${customizationArn} prefix=${customizationNamePrefix}`
195+
)
196+
return
197+
}
198+
// Since we selected based on a match, we'll reuse the persisted values.
199+
await selectCustomization(match)
185200
}
186-
// Since we selected based on a match, we'll reuse the persisted values.
187-
await selectCustomization(match)
188-
}
189201
)
190202

191203
export const showLearnMore = Commands.declare(

‎packages/core/src/test/credentials/auth.test.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import * as sinon from 'sinon'
88
import { ToolkitError, isUserCancelledError } from '../../shared/errors'
99
import { assertTreeItem } from '../shared/treeview/testUtil'
1010
import { getTestWindow } from '../shared/vscode/window'
11-
import { assertTelemetry, captureEventOnce } from '../testUtil'
11+
import { assertTelemetry, captureEventOnce, getMetrics } from '../testUtil'
1212
import { createBuilderIdProfile, createSsoProfile, createTestAuth } from './testUtil'
1313
import { toCollection } from '../../shared/utilities/asyncCollection'
1414
import globals from '../../shared/extensionGlobals'
@@ -257,7 +257,12 @@ describe('Auth', function () {
257257
it('reauthentication is indicated in metric', async function () {
258258
const conn = await auth.createInvalidSsoConnection(ssoProfile)
259259
await auth.reauthenticate(conn)
260-
assertTelemetry('aws_loginWithBrowser', { isReAuth: true })
260+
assertTelemetry('aws_loginWithBrowser', {
261+
result: 'Succeeded',
262+
isReAuth: true,
263+
credentialStartUrl: ssoProfile.startUrl,
264+
})
265+
assert.strictEqual(getMetrics('aws_loginWithBrowser').length, 1)
261266
})
262267
})
263268

‎packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,11 @@ describe('SsoAccessTokenProvider', function () {
221221
const cachedToken = await cache.token.load(startUrl).then(a => a?.token)
222222
assert.deepStrictEqual(cachedToken, token)
223223
assert.deepStrictEqual(await cache.registration.load({ region }), registration)
224-
assertTelemetry('aws_loginWithBrowser', { result: 'Succeeded', isReAuth: undefined })
224+
assertTelemetry('aws_loginWithBrowser', {
225+
result: 'Succeeded',
226+
isReAuth: undefined,
227+
credentialStartUrl: startUrl,
228+
})
225229
})
226230

227231
it('always creates a new token, even if already cached', async function () {
@@ -259,7 +263,11 @@ describe('SsoAccessTokenProvider', function () {
259263
await clock.tickAsync(750)
260264
assert.ok(!progress.visible)
261265
await assert.rejects(resp, ToolkitError)
262-
assertTelemetry('aws_loginWithBrowser', { result: 'Failed', isReAuth: undefined })
266+
assertTelemetry('aws_loginWithBrowser', {
267+
result: 'Failed',
268+
isReAuth: undefined,
269+
credentialStartUrl: startUrl,
270+
})
263271
})
264272

265273
/**
@@ -328,7 +336,11 @@ describe('SsoAccessTokenProvider', function () {
328336

329337
await assert.rejects(sut.createToken(), exception)
330338
assert.strictEqual(await cache.registration.load({ region }), undefined)
331-
assertTelemetry('aws_loginWithBrowser', { result: 'Failed', isReAuth: undefined })
339+
assertTelemetry('aws_loginWithBrowser', {
340+
result: 'Failed',
341+
isReAuth: undefined,
342+
credentialStartUrl: startUrl,
343+
})
332344
})
333345

334346
it('preserves the client registration cache on server faults', async function () {
@@ -353,7 +365,11 @@ describe('SsoAccessTokenProvider', function () {
353365
it('stops the flow if user does not click the link', async function () {
354366
stubOpen(false)
355367
await assert.rejects(sut.createToken(), ToolkitError)
356-
assertTelemetry('aws_loginWithBrowser', { result: 'Cancelled', isReAuth: undefined })
368+
assertTelemetry('aws_loginWithBrowser', {
369+
result: 'Cancelled',
370+
isReAuth: undefined,
371+
credentialStartUrl: startUrl,
372+
})
357373
})
358374

359375
it('saves the client registration even when cancelled', async function () {

0 commit comments

Comments
 (0)