Skip to content

Commit

Permalink
feat(notifications): setup and activation
Browse files Browse the repository at this point in the history
- Gated behind dev setting (`aws.dev.notifications`)
- Setup the notifications panel and begin polling in each of the extensions.
- Currently the code collects a bunch of information about the state of auth for telemetry (auth_userState). Refactor this code in each extension so that we can re-use that information for the notification rule engine.
  • Loading branch information
hayemaxi committed Nov 5, 2024
1 parent 77bd8a7 commit b4e0762
Show file tree
Hide file tree
Showing 14 changed files with 208 additions and 124 deletions.
89 changes: 42 additions & 47 deletions packages/amazonq/src/extension.ts
Original file line number Diff line number Diff line change
Expand Up @@ -46,13 +46,14 @@ import {
maybeShowMinVscodeWarning,
isSageMaker,
} from 'aws-core-vscode/shared'
import { ExtStartUpSources, telemetry } from 'aws-core-vscode/telemetry'
import { AuthUserState, ExtStartUpSources, telemetry } from 'aws-core-vscode/telemetry'
import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
import { join } from 'path'
import * as semver from 'semver'
import * as vscode from 'vscode'
import { registerCommands } from './commands'
import { focusAmazonQPanel } from 'aws-core-vscode/codewhispererChat'
import { activate as activateNotifications } from 'aws-core-vscode/notifications'

export const amazonQContextPrefix = 'amazonq'

Expand Down Expand Up @@ -162,52 +163,46 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
}, 1000)
}

await telemetry.auth_userState
.run(async () => {
telemetry.record({ passive: true })

const firstUse = AuthUtils.ExtensionUse.instance.isFirstUse()
const wasUpdated = AuthUtils.ExtensionUse.instance.wasUpdated()

if (firstUse) {
telemetry.record({ source: ExtStartUpSources.firstStartUp })
} else if (wasUpdated) {
telemetry.record({ source: ExtStartUpSources.update })
} else {
telemetry.record({ source: ExtStartUpSources.reload })
}

let authState: AuthState = 'disconnected'
try {
// May call connection validate functions that try to refresh the token.
// This could result in network errors.
authState = (await AuthUtil.instance.getChatAuthState(false)).codewhispererChat
} catch (err) {
if (
isNetworkError(err) &&
AuthUtil.instance.conn &&
AuthUtil.instance.auth.getConnectionState(AuthUtil.instance.conn) === 'valid'
) {
authState = 'connectedWithNetworkError'
} else {
throw err
}
}
const currConn = AuthUtil.instance.conn
if (currConn !== undefined && !(isAnySsoConnection(currConn) || isSageMaker())) {
getLogger().error(`Current Amazon Q connection is not SSO, type is: %s`, currConn?.type)
}

telemetry.record({
authStatus:
authState === 'connected' || authState === 'expired' || authState === 'connectedWithNetworkError'
? authState
: 'notConnected',
authEnabledConnections: AuthUtils.getAuthFormIdsFromConnection(currConn).join(','),
...(await getTelemetryMetadataForConn(currConn)),
})
})
.catch((err) => getLogger().error('Error collecting telemetry for auth_userState: %s', err))
const authState = await getAuthState()
telemetry.auth_userState.emit({
passive: true,
source: AuthUtils.ExtensionUse.instance.sourceForTelemetry(),
...authState,
})

await activateNotifications(context, authState, getAuthState)
}

async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
let authState: AuthState = 'disconnected'
try {
// May call connection validate functions that try to refresh the token.
// This could result in network errors.
authState = (await AuthUtil.instance.getChatAuthState(false)).codewhispererChat
} catch (err) {
if (
isNetworkError(err) &&
AuthUtil.instance.conn &&
AuthUtil.instance.auth.getConnectionState(AuthUtil.instance.conn) === 'valid'
) {
authState = 'connectedWithNetworkError'
} else {
throw err
}
}
const currConn = AuthUtil.instance.conn
if (currConn !== undefined && !(isAnySsoConnection(currConn) || isSageMaker())) {
getLogger().error(`Current Amazon Q connection is not SSO, type is: %s`, currConn?.type)
}

return {
authStatus:
authState === 'connected' || authState === 'expired' || authState === 'connectedWithNetworkError'
? authState
: 'notConnected',
authEnabledConnections: AuthUtils.getAuthFormIdsFromConnection(currConn).join(','),
...(await getTelemetryMetadataForConn(currConn)),
}
}

export async function deactivateCommon() {
Expand Down
14 changes: 14 additions & 0 deletions packages/core/src/auth/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ import { EcsCredentialsProvider } from './providers/ecsCredentialsProvider'
import { EnvVarsCredentialsProvider } from './providers/envVarsCredentialsProvider'
import { showMessageWithUrl } from '../shared/utilities/messages'
import { credentialHelpUrl } from '../shared/constants'
import { ExtStartUpSource } from '../shared/telemetry/util'

// iam-only excludes Builder ID and IAM Identity Center from the list of valid connections
// TODO: Understand if "iam" should include these from the list at all
Expand Down Expand Up @@ -734,6 +735,19 @@ export class ExtensionUse {
return this.wasExtensionUpdated
}

/**
* Returns a {@link ExtStartUpSource} based on the current state of the extension.
*/
sourceForTelemetry(): ExtStartUpSource {
if (this.isFirstUse()) {
return ExtStartUpSources.firstStartUp
} else if (this.wasUpdated()) {
return ExtStartUpSources.update
} else {
return ExtStartUpSources.reload
}
}

private updateMemento(key: 'isExtensionFirstUse' | 'lastExtensionVersion', val: any) {
globals.globalState.tryUpdate(key, val)
}
Expand Down
2 changes: 2 additions & 0 deletions packages/core/src/awsexplorer/activationShared.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,4 +27,6 @@ export function registerToolView(viewNode: ToolView, context: vscode.ExtensionCo
telemetry.cdk_appExpanded.emit()
}
})

return toolView
}
55 changes: 1 addition & 54 deletions packages/core/src/extension.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ import { Commands } from './shared/vscode/commands2'
import { endpointsFileUrl, githubCreateIssueUrl, githubUrl } from './shared/constants'
import { getIdeProperties, aboutExtension, isCloud9, getDocUrl } from './shared/extensionUtilities'
import { logAndShowError, logAndShowWebviewError } from './shared/utilities/logAndShowUtils'
import { AuthStatus, telemetry } from './shared/telemetry/telemetry'
import { telemetry } from './shared/telemetry/telemetry'
import { openUrl } from './shared/utilities/vsCodeUtils'
import { activateViewsShared } from './awsexplorer/activationShared'
import fs from './shared/fs/fs'
Expand All @@ -45,11 +45,6 @@ import { UriHandler } from './shared/vscode/uriHandler'
import { disableAwsSdkWarning } from './shared/awsClientBuilder'
import { FileResourceFetcher } from './shared/resourcefetcher/fileResourceFetcher'
import { ResourceFetcher } from './shared/resourcefetcher/resourcefetcher'
import { ExtStartUpSources } from './shared/telemetry/util'
import { ExtensionUse, getAuthFormIdsFromConnection } from './auth/utils'
import { Auth } from './auth'
import { AuthFormId } from './login/webview/vue/types'
import { getTelemetryMetadataForConn, isSsoConnection } from './auth/connection'
import { registerCommands } from './commands'

// In web mode everything must be in a single file, so things like the endpoints file will not be available.
Expand Down Expand Up @@ -258,51 +253,3 @@ function wrapWithProgressForCloud9(channel: vscode.OutputChannel): (typeof vscod
})
}
}

export async function emitUserState() {
await telemetry.auth_userState.run(async () => {
telemetry.record({ passive: true })

const firstUse = ExtensionUse.instance.isFirstUse()
const wasUpdated = ExtensionUse.instance.wasUpdated()

if (firstUse) {
telemetry.record({ source: ExtStartUpSources.firstStartUp })
} else if (wasUpdated) {
telemetry.record({ source: ExtStartUpSources.update })
} else {
telemetry.record({ source: ExtStartUpSources.reload })
}

let authStatus: AuthStatus = 'notConnected'
const enabledConnections: Set<AuthFormId> = new Set()
const enabledScopes: Set<string> = new Set()
if (Auth.instance.hasConnections) {
authStatus = 'expired'
;(await Auth.instance.listConnections()).forEach((conn) => {
const state = Auth.instance.getConnectionState(conn)
if (state === 'valid') {
authStatus = 'connected'
}

getAuthFormIdsFromConnection(conn).forEach((id) => enabledConnections.add(id))
if (isSsoConnection(conn)) {
conn.scopes?.forEach((s) => enabledScopes.add(s))
}
})
}

// There may be other SSO connections in toolkit, but there is no use case for
// displaying registration info for non-active connections at this time.
const activeConn = Auth.instance.activeConnection
if (activeConn?.type === 'sso') {
telemetry.record(await getTelemetryMetadataForConn(activeConn))
}

telemetry.record({
authStatus,
authEnabledConnections: [...enabledConnections].join(','),
authScopes: [...enabledScopes].join(','),
})
})
}
51 changes: 47 additions & 4 deletions packages/core/src/extensionNode.ts
Original file line number Diff line number Diff line change
Expand Up @@ -38,24 +38,27 @@ import { activate as activateDev } from './dev/activation'
import { activate as activateApplicationComposer } from './applicationcomposer/activation'
import { activate as activateRedshift } from './awsService/redshift/activation'
import { activate as activateIamPolicyChecks } from './awsService/accessanalyzer/activation'
import { activate as activateNotifications } from './notifications/activate'
import { SchemaService } from './shared/schemas'
import { AwsResourceManager } from './dynamicResources/awsResourceManager'
import globals from './shared/extensionGlobals'
import { Experiments, Settings, showSettingsFailedMsg } from './shared/settings'
import { isReleaseVersion } from './shared/vscode/env'
import { telemetry } from './shared/telemetry/telemetry'
import { AuthStatus, AuthUserState, telemetry } from './shared/telemetry/telemetry'
import { Auth, SessionSeparationPrompt } from './auth/auth'
import { getTelemetryMetadataForConn } from './auth/connection'
import { registerSubmitFeedback } from './feedback/vue/submitFeedback'
import { activateCommon, deactivateCommon, emitUserState } from './extension'
import { activateCommon, deactivateCommon } from './extension'
import { learnMoreAmazonQCommand, qExtensionPageCommand, dismissQTree } from './amazonq/explorer/amazonQChildrenNodes'
import { AuthUtil, codeWhispererCoreScopes, isPreviousQUser } from './codewhisperer/util/authUtil'
import { installAmazonQExtension } from './codewhisperer/commands/basicCommands'
import { isExtensionInstalled, VSCODE_EXTENSION_ID } from './shared/utilities'
import { ExtensionUse, initializeCredentialsProviderManager } from './auth/utils'
import { ExtensionUse, getAuthFormIdsFromConnection, initializeCredentialsProviderManager } from './auth/utils'
import { ExtStartUpSources } from './shared/telemetry'
import { activate as activateThreatComposerEditor } from './threatComposer/activation'
import { isSsoConnection, hasScopes } from './auth/connection'
import { CrashMonitoring, setContext } from './shared'
import { AuthFormId } from './login/webview/vue/types'

let localize: nls.LocalizeFunc

Expand Down Expand Up @@ -231,7 +234,14 @@ export async function activate(context: vscode.ExtensionContext) {
globals.telemetry.assertPassiveTelemetry(globals.didReload)
}

await emitUserState()
const authState = await getAuthState()
telemetry.auth_userState.emit({
passive: true,
source: ExtensionUse.instance.sourceForTelemetry(),
...authState,
})

await activateNotifications(context, authState, getAuthState)
} catch (error) {
const stacktrace = (error as Error).stack?.split('\n')
// truncate if the stacktrace is unusually long
Expand Down Expand Up @@ -322,3 +332,36 @@ function recordToolkitInitialization(activationStartedOn: number, settingsValid:
logger?.error(err as Error)
}
}

async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
let authStatus: AuthStatus = 'notConnected'
const enabledConnections: Set<AuthFormId> = new Set()
const enabledScopes: Set<string> = new Set()
if (Auth.instance.hasConnections) {
authStatus = 'expired'
;(await Auth.instance.listConnections()).forEach((conn) => {
const state = Auth.instance.getConnectionState(conn)
if (state === 'valid') {
authStatus = 'connected'
}

getAuthFormIdsFromConnection(conn).forEach((id) => enabledConnections.add(id))
if (isSsoConnection(conn)) {
conn.scopes?.forEach((s) => enabledScopes.add(s))
}
})
}

// There may be other SSO connections in toolkit, but there is no use case for
// displaying registration info for non-active connections at this time.
const activeConn = Auth.instance.activeConnection
if (activeConn?.type === 'sso') {
telemetry.record(await getTelemetryMetadataForConn(activeConn))
}

return {
authStatus,
authEnabledConnections: [...enabledConnections].sort().join(','),
authScopes: [...enabledScopes].sort().join(','),
}
}
47 changes: 47 additions & 0 deletions packages/core/src/notifications/activate.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
/*!
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

import * as vscode from 'vscode'
import { DevSettings } from '../shared/settings'
import { NotificationsController } from './controller'
import { NotificationsNode } from './panelNode'
import { RuleEngine, getRuleContext } from './rules'
import globals from '../shared/extensionGlobals'
import { AuthState } from './types'

/** Time in MS to poll for emergency notifications */
const emergencyPollTime = 1000 * 10 * 60

/**
* Activate the in-IDE notifications module and begin receiving notifications.
*
* @param context extension context
* @param initialState initial auth state
* @param authStateFn fn to get current auth state
*/
export async function activate(
context: vscode.ExtensionContext,
initialState: AuthState,
authStateFn: () => Promise<AuthState>
) {
// TODO: Currently gated behind feature-flag.
if (!DevSettings.instance.get('notifications', false)) {
return
}

const panelNode = NotificationsNode.instance
panelNode.registerView(context)

const controller = new NotificationsController(panelNode)
const engine = new RuleEngine(await getRuleContext(context, initialState))

void controller.pollForStartUp(engine)
void controller.pollForEmergencies(engine)

globals.clock.setInterval(async () => {
const ruleContext = await getRuleContext(context, await authStateFn())
await controller.pollForEmergencies(new RuleEngine(ruleContext))
}, emergencyPollTime)
}
5 changes: 1 addition & 4 deletions packages/core/src/notifications/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,4 @@
* SPDX-License-Identifier: Apache-2.0
*/

export { RuleContext } from './types'
export { NotificationsController } from './controller'
export { RuleEngine } from './rules'
export { registerProvider, NotificationsNode } from './panelNode'
export { activate } from './activate'
Loading

0 comments on commit b4e0762

Please sign in to comment.