diff --git a/packages/amazonq/package.json b/packages/amazonq/package.json index cd0e896824c..7f7f1133cd8 100644 --- a/packages/amazonq/package.json +++ b/packages/amazonq/package.json @@ -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": [ @@ -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%", @@ -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", diff --git a/packages/core/package.nls.json b/packages/core/package.nls.json index c78dc51922c..da644ab59ef 100644 --- a/packages/core/package.nls.json +++ b/packages/core/package.nls.json @@ -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", diff --git a/packages/core/src/auth/auth.ts b/packages/core/src/auth/auth.ts index f7468a0e44b..d7c2a5d58b7 100644 --- a/packages/core/src/auth/auth.ts +++ b/packages/core/src/auth/auth.ts @@ -59,6 +59,7 @@ import { AwsConnection, scopesCodeWhispererCore, ProfileNotFoundError, + isSsoConnection, } from './connection' import { isSageMaker, isCloud9, isAmazonQ } from '../shared/extensionUtilities' import { telemetry } from '../shared/telemetry/telemetry' @@ -66,6 +67,7 @@ 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 { /** @@ -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 */ @@ -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()) + return conn } @@ -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 }) @@ -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): Promise { diff --git a/packages/core/src/codewhispererChat/commands/registerCommands.ts b/packages/core/src/codewhispererChat/commands/registerCommands.ts index 21e65bc6c83..abf6c6fa07c 100644 --- a/packages/core/src/codewhispererChat/commands/registerCommands.ts +++ b/packages/core/src/codewhispererChat/commands/registerCommands.ts @@ -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 = @@ -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' diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts index 0a9a42ec101..3926cadb3d4 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts @@ -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) { diff --git a/packages/core/src/codewhispererChat/controllers/chat/prompts/promptsGenerator.ts b/packages/core/src/codewhispererChat/controllers/chat/prompts/promptsGenerator.ts index 6abecda2f8c..d738eddcb67 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/prompts/promptsGenerator.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/prompts/promptsGenerator.ts @@ -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 { diff --git a/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts b/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts index 036f0d9e155..11ded9b27cd 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts @@ -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 } diff --git a/packages/core/src/codewhispererChat/controllers/chat/userIntent/userIntentRecognizer.ts b/packages/core/src/codewhispererChat/controllers/chat/userIntent/userIntentRecognizer.ts index e87a2ed93e6..32b6b274a51 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/userIntent/userIntentRecognizer.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/userIntent/userIntentRecognizer.ts @@ -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 { @@ -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 } @@ -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 } diff --git a/packages/core/src/shared/telemetry/vscodeTelemetry.json b/packages/core/src/shared/telemetry/vscodeTelemetry.json index eda52e0b0ed..408c96ec4fe 100644 --- a/packages/core/src/shared/telemetry/vscodeTelemetry.json +++ b/packages/core/src/shared/telemetry/vscodeTelemetry.json @@ -58,7 +58,8 @@ "showExample", "citeSources", "explainLineByLine", - "explainCodeSelection" + "explainCodeSelection", + "generateUnitTests" ], "description": "Explict user intent associated with a chat message" }, diff --git a/packages/core/src/shared/vscode/setContext.ts b/packages/core/src/shared/vscode/setContext.ts index 9cf7c33c25f..5068bc59e13 100644 --- a/packages/core/src/shared/vscode/setContext.ts +++ b/packages/core/src/shared/vscode/setContext.ts @@ -14,6 +14,7 @@ type contextKey = | 'aws.isDevMode' | 'aws.isSageMaker' | 'aws.isWebExtHost' + | 'aws.isInternalUser' | 'aws.amazonq.showLoginView' | 'aws.codecatalyst.connected' | 'aws.codewhisperer.connected' diff --git a/src.gen/@amzn/codewhisperer-streaming/src/commands/GenerateAssistantResponseCommand.ts b/src.gen/@amzn/codewhisperer-streaming/src/commands/GenerateAssistantResponseCommand.ts index 769ecff549f..87e0f6fa51f 100644 --- a/src.gen/@amzn/codewhisperer-streaming/src/commands/GenerateAssistantResponseCommand.ts +++ b/src.gen/@amzn/codewhisperer-streaming/src/commands/GenerateAssistantResponseCommand.ts @@ -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", @@ -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", * }, * }, * }, @@ -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", @@ -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", * }, * }, * }, @@ -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 diff --git a/src.gen/@amzn/codewhisperer-streaming/src/commands/GenerateTaskAssistPlanCommand.ts b/src.gen/@amzn/codewhisperer-streaming/src/commands/GenerateTaskAssistPlanCommand.ts index fc884f5f339..22a591f22d8 100644 --- a/src.gen/@amzn/codewhisperer-streaming/src/commands/GenerateTaskAssistPlanCommand.ts +++ b/src.gen/@amzn/codewhisperer-streaming/src/commands/GenerateTaskAssistPlanCommand.ts @@ -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", @@ -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", * }, * }, * }, @@ -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", @@ -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", * }, * }, * }, @@ -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 diff --git a/src.gen/@amzn/codewhisperer-streaming/src/models/models_0.ts b/src.gen/@amzn/codewhisperer-streaming/src/models/models_0.ts index fa424d76496..7d2622c985e 100644 --- a/src.gen/@amzn/codewhisperer-streaming/src/models/models_0.ts +++ b/src.gen/@amzn/codewhisperer-streaming/src/models/models_0.ts @@ -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 */