Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
a16fd11
feat (Amazon Q): Add basic UX and logic for generating unit tests
bweedop Aug 22, 2024
fcc3818
Merge branch 'master' into add-generate-tests
bweedop Sep 9, 2024
5cc1126
Add GENERATE_UNIT_TESTS to model; Limit access to test command to int…
bweedop Sep 11, 2024
f2ff8f7
Correct package.json
bweedop Sep 11, 2024
50fc1f8
Add Beta suffix to generate unit tests command
bweedop Sep 11, 2024
3ed55a9
Merge branch 'aws:master' into add-generate-tests
bweedop Sep 11, 2024
334dd2a
Refactor to reduce code duplication
bweedop Sep 11, 2024
3e3b6aa
Merge branch 'add-generate-tests' of github.com:bweedop/aws-toolkit-v…
bweedop Sep 11, 2024
4b19d4a
Revert removal of suppress prompts
bweedop Sep 11, 2024
4dc4f6b
Remove unnecessary imports
bweedop Sep 11, 2024
ab799b5
Merge branch 'aws:master' into add-generate-tests
bweedop Sep 12, 2024
ecc91a1
Merge branch 'aws:master' into add-generate-tests
bweedop Sep 17, 2024
15de7e9
Merge branch 'master' into add-generate-tests
bweedop Sep 17, 2024
17a9153
Resolve circular dependency issue
bweedop Sep 21, 2024
5701cef
Add generateUnitTests to telemetry
bweedop Sep 23, 2024
a2c7dc3
Move check for internal Amazon user to general auth
bweedop Sep 23, 2024
99cd07d
Merge branch 'aws:master' into add-generate-tests
bweedop Sep 23, 2024
59a023b
Update naming of reformatStartUrl method
bweedop Sep 24, 2024
9a11570
Move comment to docstring
bweedop Sep 24, 2024
a12768c
set context on forgotten and expired connection
bweedop Sep 24, 2024
12f879d
Default isInternalUser to false on forgotten and expired connection
bweedop Sep 25, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 19 additions & 1 deletion packages/amazonq/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -349,9 +349,14 @@
"command": "aws.amazonq.optimizeCode",
"group": "cw_chat@4"
},
{
"command": "aws.amazonq.generateUnitTests",
"group": "cw_chat@5",
"when": "aws.codewhisperer.connected && aws.isInternalUser"
},
{
"command": "aws.amazonq.sendToPrompt",
"group": "cw_chat@5"
"group": "cw_chat@6"
}
],
"editor/context": [
Expand Down Expand Up @@ -423,6 +428,12 @@
"category": "%AWS.amazonq.title%",
"enablement": "aws.codewhisperer.connected"
},
{
"command": "aws.amazonq.generateUnitTests",
"title": "%AWS.command.amazonq.generateUnitTests%",
"category": "%AWS.amazonq.title%",
"enablement": "aws.codewhisperer.connected && aws.isInternalUser"
},
{
"command": "aws.amazonq.reconnect",
"title": "%AWS.command.codewhisperer.reconnect%",
Expand Down Expand Up @@ -599,6 +610,13 @@
"mac": "cmd+alt+q",
"linux": "meta+alt+q"
},
{
"command": "aws.amazonq.generateUnitTests",
"key": "win+alt+t",
"mac": "cmd+alt+t",
"linux": "meta+alt+t",
"when": "aws.codewhisperer.connected && aws.isInternalUser"
},
{
"command": "aws.amazonq.invokeInlineCompletion",
"key": "alt+c",
Expand Down
1 change: 1 addition & 0 deletions packages/core/package.nls.json
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,7 @@
"AWS.command.amazonq.fixCode": "Fix",
"AWS.command.amazonq.optimizeCode": "Optimize",
"AWS.command.amazonq.sendToPrompt": "Send to prompt",
"AWS.command.amazonq.generateUnitTests": "Generate Tests (Beta)",
"AWS.command.amazonq.security.scan": "Run Project Scan",
"AWS.command.deploySamApplication": "Deploy SAM Application",
"AWS.command.aboutToolkit": "About",
Expand Down
30 changes: 30 additions & 0 deletions packages/core/src/auth/auth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -59,13 +59,15 @@ import {
AwsConnection,
scopesCodeWhispererCore,
ProfileNotFoundError,
isSsoConnection,
} from './connection'
import { isSageMaker, isCloud9, isAmazonQ } from '../shared/extensionUtilities'
import { telemetry } from '../shared/telemetry/telemetry'
import { randomUUID } from '../shared/crypto'
import { asStringifiedStack } from '../shared/telemetry/spans'
import { withTelemetryContext } from '../shared/telemetry/util'
import { DiskCacheError } from '../shared/utilities/cacheUtils'
import { setContext } from '../shared/vscode/setContext'

interface AuthService {
/**
Expand Down Expand Up @@ -166,6 +168,30 @@ export class Auth implements AuthService, ConnectionManager {
return this.#ssoCacheWatcher
}

public get startUrl(): string | undefined {
return isSsoConnection(this.activeConnection)
? this.normalizeStartUrl(this.activeConnection.startUrl)
: undefined
}

public isConnected(): boolean {
return this.activeConnection !== undefined
}

/**
* Normalizes the provided URL
*
* Any trailing '/' and `#` is removed from the URL
* e.g. https://view.awsapps.com/start/# will become https://view.awsapps.com/start
*/
public normalizeStartUrl(startUrl: string | undefined) {
return !startUrl ? undefined : startUrl.replace(/[\/#]+$/g, '')
}

public isInternalAmazonUser(): boolean {
return this.isConnected() && this.startUrl === 'https://amzn.awsapps.com/start'
}

/**
* Map startUrl -> declared connections
*/
Expand Down Expand Up @@ -223,6 +249,8 @@ export class Auth implements AuthService, ConnectionManager {
this.#onDidChangeActiveConnection.fire(conn)
await this.store.setCurrentProfileId(id)

await setContext('aws.isInternalUser', this.isInternalAmazonUser())
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was wondering if forgetConnection and expireConnection should also un-set this. But it's probably useful to keep this hint around, once it was deduced. And in that case, should this flag be stored in globalState? Then we don't need to wait until a valid connection next time.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think it hurts to set this back to false on forgetConnection and expireConnection.

Some testing that I did made me doubt how globalState worked. Using globalState.update for this didn't work to enable the feature for internal users. setContext, however, worked. Are global keys detectable from package.json after being set?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are global keys detectable from package.json after being set?

No

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh okay well that would explain it. We need this to be set with setContext


return conn
}

Expand Down Expand Up @@ -373,6 +401,7 @@ export class Auth implements AuthService, ConnectionManager {
}
}
this.#onDidDeleteConnection.fire({ connId, storedProfile: profile })
await setContext('aws.isInternalUser', false)
}

@withTelemetryContext({ name: 'clearStaleLinkedIamConnections', class: authClassName })
Expand Down Expand Up @@ -405,6 +434,7 @@ export class Auth implements AuthService, ConnectionManager {
await provider.invalidate('devModeManualExpiration')
// updates the state of the connection
await this.refreshConnectionState(conn)
await setContext('aws.isInternalUser', false)
}

public async getConnection(connection: Pick<Connection, 'id'>): Promise<Connection | undefined> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,14 @@ export function registerCommands(controllerPublishers: ChatControllerMessagePubl
})
})
})
Commands.register('aws.amazonq.generateUnitTests', async (data) => {
return focusAmazonQPanel.execute(placeholder, 'amazonq.generateUnitTests').then(() => {
controllerPublishers.processContextMenuCommand.publish({
type: 'aws.amazonq.generateUnitTests',
triggerType: getCommandTriggerType(data),
})
})
})
}

export type EditorContextBaseCommandType =
Expand All @@ -102,6 +110,7 @@ export type EditorContextBaseCommandType =
| 'aws.amazonq.fixCode'
| 'aws.amazonq.optimizeCode'
| 'aws.amazonq.sendToPrompt'
| 'aws.amazonq.generateUnitTests'

export type CodeScanIssueCommandType = 'aws.amazonq.explainIssue'

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -339,6 +339,7 @@ export class Messenger {
['aws.amazonq.fixCode', 'Fix'],
['aws.amazonq.optimizeCode', 'Optimize'],
['aws.amazonq.sendToPrompt', 'Send to prompt'],
['aws.amazonq.generateUnitTests', 'Generate unit tests for'],
])

public sendStaticTextResponse(type: StaticTextResponseType, triggerID: string, tabID: string) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ export class PromptsGenerator {
['aws.amazonq.fixCode', 'Fix'],
['aws.amazonq.optimizeCode', 'Optimize'],
['aws.amazonq.sendToPrompt', 'Send to prompt'],
['aws.amazonq.generateUnitTests', 'Generate unit tests for'],
])

public generateForContextMenuCommand(command: EditorContextCommand): string {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,8 @@ export class CWCTelemetryHelper {
return 'explainLineByLine'
case UserIntent.SHOW_EXAMPLES:
return 'showExample'
case UserIntent.GENERATE_UNIT_TESTS:
return 'generateUnitTests'
default:
return undefined
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import { UserIntent } from '@amzn/codewhisperer-streaming'
import { EditorContextCommand } from '../../../commands/registerCommands'
import { PromptMessage } from '../model'
import { Auth } from '../../../../auth'

export class UserIntentRecognizer {
public getFromContextMenuCommand(command: EditorContextCommand): UserIntent | undefined {
Expand All @@ -18,6 +19,8 @@ export class UserIntentRecognizer {
return UserIntent.APPLY_COMMON_BEST_PRACTICES
case 'aws.amazonq.optimizeCode':
return UserIntent.IMPROVE_CODE
case 'aws.amazonq.generateUnitTests':
return UserIntent.GENERATE_UNIT_TESTS
default:
return undefined
}
Expand All @@ -36,6 +39,8 @@ export class UserIntentRecognizer {
return UserIntent.APPLY_COMMON_BEST_PRACTICES
} else if (prompt.message.startsWith('Optimize')) {
return UserIntent.IMPROVE_CODE
} else if (prompt.message.startsWith('Generate unit tests') && Auth.instance.isInternalAmazonUser()) {
return UserIntent.GENERATE_UNIT_TESTS
}
return undefined
}
Expand Down
3 changes: 2 additions & 1 deletion packages/core/src/shared/telemetry/vscodeTelemetry.json
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,8 @@
"showExample",
"citeSources",
"explainLineByLine",
"explainCodeSelection"
"explainCodeSelection",
"generateUnitTests"
],
"description": "Explict user intent associated with a chat message"
},
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/shared/vscode/setContext.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ type contextKey =
| 'aws.isDevMode'
| 'aws.isSageMaker'
| 'aws.isWebExtHost'
| 'aws.isInternalUser'
| 'aws.amazonq.showLoginView'
| 'aws.codecatalyst.connected'
| 'aws.codewhisperer.connected'
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,7 @@ export interface GenerateAssistantResponseCommandOutput extends GenerateAssistan
* hasConsentedToCrossRegionCalls: true || false,
* },
* },
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE",
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE" || "GENERATE_UNIT_TESTS",
* },
* assistantResponseMessage: { // AssistantResponseMessage
* messageId: "STRING_VALUE",
Expand All @@ -214,7 +214,7 @@ export interface GenerateAssistantResponseCommandOutput extends GenerateAssistan
* ],
* followupPrompt: { // FollowupPrompt
* content: "STRING_VALUE", // required
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE",
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE" || "GENERATE_UNIT_TESTS",
* },
* },
* },
Expand Down Expand Up @@ -305,7 +305,7 @@ export interface GenerateAssistantResponseCommandOutput extends GenerateAssistan
* hasConsentedToCrossRegionCalls: true || false,
* },
* },
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE",
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE" || "GENERATE_UNIT_TESTS",
* },
* assistantResponseMessage: {
* messageId: "STRING_VALUE",
Expand All @@ -330,7 +330,7 @@ export interface GenerateAssistantResponseCommandOutput extends GenerateAssistan
* ],
* followupPrompt: {
* content: "STRING_VALUE", // required
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE",
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE" || "GENERATE_UNIT_TESTS",
* },
* },
* },
Expand Down Expand Up @@ -375,7 +375,7 @@ export interface GenerateAssistantResponseCommandOutput extends GenerateAssistan
* // followupPromptEvent: { // FollowupPromptEvent
* // followupPrompt: { // FollowupPrompt
* // content: "STRING_VALUE", // required
* // userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE",
* // userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE" || "GENERATE_UNIT_TESTS",
* // },
* // },
* // codeEvent: { // CodeEvent
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,7 @@ export interface GenerateTaskAssistPlanCommandOutput extends GenerateTaskAssistP
* hasConsentedToCrossRegionCalls: true || false,
* },
* },
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE",
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE" || "GENERATE_UNIT_TESTS",
* },
* assistantResponseMessage: { // AssistantResponseMessage
* messageId: "STRING_VALUE",
Expand All @@ -214,7 +214,7 @@ export interface GenerateTaskAssistPlanCommandOutput extends GenerateTaskAssistP
* ],
* followupPrompt: { // FollowupPrompt
* content: "STRING_VALUE", // required
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE",
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE" || "GENERATE_UNIT_TESTS",
* },
* },
* },
Expand Down Expand Up @@ -305,7 +305,7 @@ export interface GenerateTaskAssistPlanCommandOutput extends GenerateTaskAssistP
* hasConsentedToCrossRegionCalls: true || false,
* },
* },
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE",
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE" || "GENERATE_UNIT_TESTS",
* },
* assistantResponseMessage: {
* messageId: "STRING_VALUE",
Expand All @@ -330,7 +330,7 @@ export interface GenerateTaskAssistPlanCommandOutput extends GenerateTaskAssistP
* ],
* followupPrompt: {
* content: "STRING_VALUE", // required
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE",
* userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE" || "GENERATE_UNIT_TESTS",
* },
* },
* },
Expand Down Expand Up @@ -378,7 +378,7 @@ export interface GenerateTaskAssistPlanCommandOutput extends GenerateTaskAssistP
* // followupPromptEvent: { // FollowupPromptEvent
* // followupPrompt: { // FollowupPrompt
* // content: "STRING_VALUE", // required
* // userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE",
* // userIntent: "SUGGEST_ALTERNATE_IMPLEMENTATION" || "APPLY_COMMON_BEST_PRACTICES" || "IMPROVE_CODE" || "SHOW_EXAMPLES" || "CITE_SOURCES" || "EXPLAIN_LINE_BY_LINE" || "EXPLAIN_CODE_SELECTION" || "GENERATE_CLOUDFORMATION_TEMPLATE" || "GENERATE_UNIT_TESTS",
* // },
* // },
* // codeEvent: { // CodeEvent
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -245,6 +245,10 @@ export const UserIntent = {
* Generate CloudFormation Template
*/
GENERATE_CLOUDFORMATION_TEMPLATE: "GENERATE_CLOUDFORMATION_TEMPLATE",
/**
* Generate Unit Tests
*/
GENERATE_UNIT_TESTS: "GENERATE_UNIT_TESTS",
/**
* Improve Code
*/
Expand Down
Loading