diff --git a/docs/lsp.md b/docs/lsp.md index 42d94d334a4..a0c7a25d8cb 100644 --- a/docs/lsp.md +++ b/docs/lsp.md @@ -26,9 +26,7 @@ sequenceDiagram ## Language Server Debugging -1. Clone https://github.com/aws/language-servers.git and set it up in the same workspace as this project by cmd+shift+p and "add folder to workspace" and selecting the language-servers folder that you just cloned. Your VS code folder structure should look like below. - - +1. Clone https://github.com/aws/language-servers.git and set it up in the same workspace as this project by cmd+shift+p and "add folder to workspace" and selecting the language-servers folder that you just cloned. Your VS code folder structure should look like below. ``` /aws-toolkit-vscode @@ -48,7 +46,6 @@ sequenceDiagram 3. Enable the lsp experiment: ``` "aws.experiments": { - "amazonqLSP": true, "amazonqLSPInline": true, // optional: enables inline completion from flare "amazonqLSPChat": true // optional: enables chat from flare } diff --git a/package-lock.json b/package-lock.json index aa538ee968b..51483c52c1e 100644 --- a/package-lock.json +++ b/package-lock.json @@ -10880,8 +10880,6 @@ }, "node_modules/@aws-toolkits/telemetry": { "version": "1.0.322", - "resolved": "https://registry.npmjs.org/@aws-toolkits/telemetry/-/telemetry-1.0.322.tgz", - "integrity": "sha512-KtLabV3ycRH31EAZ0xoWrdpIBG3ym8CQAqgkHd9DSefndbepPRa07atfXw73Ok9J5aA81VHCFpx1dwrLg39EcQ==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -10906,8 +10904,6 @@ }, "node_modules/@aws/chat-client-ui-types": { "version": "0.1.26", - "resolved": "https://registry.npmjs.org/@aws/chat-client-ui-types/-/chat-client-ui-types-0.1.26.tgz", - "integrity": "sha512-WlF0fP1nojueknr815dg6Ivs+Q3e5onvWTH1nI05jysSzUHjsWwFDBrsxqJXfaPIFhPrbQzHqoxHbhIwQ1OLuw==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -10916,8 +10912,6 @@ }, "node_modules/@aws/language-server-runtimes": { "version": "0.2.81", - "resolved": "https://registry.npmjs.org/@aws/language-server-runtimes/-/language-server-runtimes-0.2.81.tgz", - "integrity": "sha512-wnwa8ctVCAckIpfWSblHyLVzl6UKX5G7ft+yetH1pI0mZvseSNzHUhclxNl4WGaDgGnEbBjLD0XRNEy2yRrSYg==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -10946,8 +10940,6 @@ }, "node_modules/@aws/language-server-runtimes-types": { "version": "0.1.28", - "resolved": "https://registry.npmjs.org/@aws/language-server-runtimes-types/-/language-server-runtimes-types-0.1.28.tgz", - "integrity": "sha512-eDNcEXGAyD4rzl+eVJ6Ngfbm4iaR8MkoMk1wVcnV+VGqu63TyvV1aVWnZdl9tR4pmC0rIH3tj8FSCjhSU6eJlA==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -10957,8 +10949,6 @@ }, "node_modules/@aws/language-server-runtimes/node_modules/@smithy/abort-controller": { "version": "4.0.2", - "resolved": "https://registry.npmjs.org/@smithy/abort-controller/-/abort-controller-4.0.2.tgz", - "integrity": "sha512-Sl/78VDtgqKxN2+1qduaVE140XF+Xg+TafkncspwM4jFP/LHr76ZHmIY/y3V1M0mMLNk+Je6IGbzxy23RSToMw==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -10971,8 +10961,6 @@ }, "node_modules/@aws/language-server-runtimes/node_modules/@smithy/node-http-handler": { "version": "4.0.4", - "resolved": "https://registry.npmjs.org/@smithy/node-http-handler/-/node-http-handler-4.0.4.tgz", - "integrity": "sha512-/mdqabuAT3o/ihBGjL94PUbTSPSRJ0eeVTdgADzow0wRJ0rN4A27EOrtlK56MYiO1fDvlO3jVTCxQtQmK9dZ1g==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -10988,8 +10976,6 @@ }, "node_modules/@aws/language-server-runtimes/node_modules/@smithy/protocol-http": { "version": "5.1.0", - "resolved": "https://registry.npmjs.org/@smithy/protocol-http/-/protocol-http-5.1.0.tgz", - "integrity": "sha512-KxAOL1nUNw2JTYrtviRRjEnykIDhxc84qMBzxvu1MUfQfHTuBlCG7PA6EdVwqpJjH7glw7FqQoFxUJSyBQgu7g==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -11002,8 +10988,6 @@ }, "node_modules/@aws/language-server-runtimes/node_modules/@smithy/querystring-builder": { "version": "4.0.2", - "resolved": "https://registry.npmjs.org/@smithy/querystring-builder/-/querystring-builder-4.0.2.tgz", - "integrity": "sha512-NTOs0FwHw1vimmQM4ebh+wFQvOwkEf/kQL6bSM1Lock+Bv4I89B3hGYoUEPkmvYPkDKyp5UdXJYu+PoTQ3T31Q==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -11017,8 +11001,6 @@ }, "node_modules/@aws/language-server-runtimes/node_modules/@smithy/types": { "version": "4.2.0", - "resolved": "https://registry.npmjs.org/@smithy/types/-/types-4.2.0.tgz", - "integrity": "sha512-7eMk09zQKCO+E/ivsjQv+fDlOupcFUCSC/L2YUPgwhvowVGWbPQHjEFcmjt7QQ4ra5lyowS92SV53Zc6XD4+fg==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -11030,8 +11012,6 @@ }, "node_modules/@aws/language-server-runtimes/node_modules/@smithy/util-uri-escape": { "version": "4.0.0", - "resolved": "https://registry.npmjs.org/@smithy/util-uri-escape/-/util-uri-escape-4.0.0.tgz", - "integrity": "sha512-77yfbCbQMtgtTylO9itEAdpPXSog3ZxMe09AEhm0dU0NLTalV70ghDZFR+Nfi1C60jnJoh/Re4090/DuZh2Omg==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -11099,8 +11079,6 @@ }, "node_modules/@aws/mynah-ui": { "version": "4.30.3", - "resolved": "https://registry.npmjs.org/@aws/mynah-ui/-/mynah-ui-4.30.3.tgz", - "integrity": "sha512-Xy22dzCaFUqpdSHMpLa8Dsq98DiAUq49dm7Iu8Yj2YZXSCyfKQiYMJOfwU8IoqeNcEney5JRMJpf+/RysWugbA==", "hasInstallScript": true, "license": "Apache License 2.0", "dependencies": { @@ -11674,8 +11652,6 @@ }, "node_modules/@opentelemetry/api": { "version": "1.9.0", - "resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.0.tgz", - "integrity": "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg==", "dev": true, "license": "Apache-2.0", "engines": { @@ -11684,8 +11660,6 @@ }, "node_modules/@opentelemetry/api-logs": { "version": "0.200.0", - "resolved": "https://registry.npmjs.org/@opentelemetry/api-logs/-/api-logs-0.200.0.tgz", - "integrity": "sha512-IKJBQxh91qJ+3ssRly5hYEJ8NDHu9oY/B1PXVSCWf7zytmYO9RNLB0Ox9XQ/fJ8m6gY6Q6NtBWlmXfaXt5Uc4Q==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -11696,9 +11670,7 @@ } }, "node_modules/@opentelemetry/core": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/@opentelemetry/core/-/core-2.0.0.tgz", - "integrity": "sha512-SLX36allrcnVaPYG3R78F/UZZsBsvbc7lMCLx37LyH5MJ1KAAZ2E3mW9OAD3zGz0G8q/BtoS5VUrjzDydhD6LQ==", + "version": "2.0.1", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -11713,8 +11685,6 @@ }, "node_modules/@opentelemetry/exporter-logs-otlp-http": { "version": "0.200.0", - "resolved": "https://registry.npmjs.org/@opentelemetry/exporter-logs-otlp-http/-/exporter-logs-otlp-http-0.200.0.tgz", - "integrity": "sha512-KfWw49htbGGp9s8N4KI8EQ9XuqKJ0VG+yVYVYFiCYSjEV32qpQ5qZ9UZBzOZ6xRb+E16SXOSCT3RkqBVSABZ+g==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -11731,10 +11701,22 @@ "@opentelemetry/api": "^1.3.0" } }, + "node_modules/@opentelemetry/exporter-logs-otlp-http/node_modules/@opentelemetry/core": { + "version": "2.0.0", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@opentelemetry/semantic-conventions": "^1.29.0" + }, + "engines": { + "node": "^18.19.0 || >=20.6.0" + }, + "peerDependencies": { + "@opentelemetry/api": ">=1.0.0 <1.10.0" + } + }, "node_modules/@opentelemetry/exporter-metrics-otlp-http": { "version": "0.200.0", - "resolved": "https://registry.npmjs.org/@opentelemetry/exporter-metrics-otlp-http/-/exporter-metrics-otlp-http-0.200.0.tgz", - "integrity": "sha512-5BiR6i8yHc9+qW7F6LqkuUnIzVNA7lt0qRxIKcKT+gq3eGUPHZ3DY29sfxI3tkvnwMgtnHDMNze5DdxW39HsAw==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -11751,10 +11733,52 @@ "@opentelemetry/api": "^1.3.0" } }, + "node_modules/@opentelemetry/exporter-metrics-otlp-http/node_modules/@opentelemetry/core": { + "version": "2.0.0", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@opentelemetry/semantic-conventions": "^1.29.0" + }, + "engines": { + "node": "^18.19.0 || >=20.6.0" + }, + "peerDependencies": { + "@opentelemetry/api": ">=1.0.0 <1.10.0" + } + }, + "node_modules/@opentelemetry/exporter-metrics-otlp-http/node_modules/@opentelemetry/resources": { + "version": "2.0.0", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@opentelemetry/core": "2.0.0", + "@opentelemetry/semantic-conventions": "^1.29.0" + }, + "engines": { + "node": "^18.19.0 || >=20.6.0" + }, + "peerDependencies": { + "@opentelemetry/api": ">=1.3.0 <1.10.0" + } + }, + "node_modules/@opentelemetry/exporter-metrics-otlp-http/node_modules/@opentelemetry/sdk-metrics": { + "version": "2.0.0", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@opentelemetry/core": "2.0.0", + "@opentelemetry/resources": "2.0.0" + }, + "engines": { + "node": "^18.19.0 || >=20.6.0" + }, + "peerDependencies": { + "@opentelemetry/api": ">=1.9.0 <1.10.0" + } + }, "node_modules/@opentelemetry/otlp-exporter-base": { "version": "0.200.0", - "resolved": "https://registry.npmjs.org/@opentelemetry/otlp-exporter-base/-/otlp-exporter-base-0.200.0.tgz", - "integrity": "sha512-IxJgA3FD7q4V6gGq4bnmQM5nTIyMDkoGFGrBrrDjB6onEiq1pafma55V+bHvGYLWvcqbBbRfezr1GED88lacEQ==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -11768,10 +11792,22 @@ "@opentelemetry/api": "^1.3.0" } }, + "node_modules/@opentelemetry/otlp-exporter-base/node_modules/@opentelemetry/core": { + "version": "2.0.0", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@opentelemetry/semantic-conventions": "^1.29.0" + }, + "engines": { + "node": "^18.19.0 || >=20.6.0" + }, + "peerDependencies": { + "@opentelemetry/api": ">=1.0.0 <1.10.0" + } + }, "node_modules/@opentelemetry/otlp-transformer": { "version": "0.200.0", - "resolved": "https://registry.npmjs.org/@opentelemetry/otlp-transformer/-/otlp-transformer-0.200.0.tgz", - "integrity": "sha512-+9YDZbYybOnv7sWzebWOeK6gKyt2XE7iarSyBFkwwnP559pEevKOUD8NyDHhRjCSp13ybh9iVXlMfcj/DwF/yw==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -11790,14 +11826,56 @@ "@opentelemetry/api": "^1.3.0" } }, - "node_modules/@opentelemetry/resources": { + "node_modules/@opentelemetry/otlp-transformer/node_modules/@opentelemetry/core": { + "version": "2.0.0", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@opentelemetry/semantic-conventions": "^1.29.0" + }, + "engines": { + "node": "^18.19.0 || >=20.6.0" + }, + "peerDependencies": { + "@opentelemetry/api": ">=1.0.0 <1.10.0" + } + }, + "node_modules/@opentelemetry/otlp-transformer/node_modules/@opentelemetry/resources": { + "version": "2.0.0", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@opentelemetry/core": "2.0.0", + "@opentelemetry/semantic-conventions": "^1.29.0" + }, + "engines": { + "node": "^18.19.0 || >=20.6.0" + }, + "peerDependencies": { + "@opentelemetry/api": ">=1.3.0 <1.10.0" + } + }, + "node_modules/@opentelemetry/otlp-transformer/node_modules/@opentelemetry/sdk-metrics": { "version": "2.0.0", - "resolved": "https://registry.npmjs.org/@opentelemetry/resources/-/resources-2.0.0.tgz", - "integrity": "sha512-rnZr6dML2z4IARI4zPGQV4arDikF/9OXZQzrC01dLmn0CZxU5U5OLd/m1T7YkGRj5UitjeoCtg/zorlgMQcdTg==", "dev": true, "license": "Apache-2.0", "dependencies": { "@opentelemetry/core": "2.0.0", + "@opentelemetry/resources": "2.0.0" + }, + "engines": { + "node": "^18.19.0 || >=20.6.0" + }, + "peerDependencies": { + "@opentelemetry/api": ">=1.9.0 <1.10.0" + } + }, + "node_modules/@opentelemetry/resources": { + "version": "2.0.1", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@opentelemetry/core": "2.0.1", "@opentelemetry/semantic-conventions": "^1.29.0" }, "engines": { @@ -11809,8 +11887,6 @@ }, "node_modules/@opentelemetry/sdk-logs": { "version": "0.200.0", - "resolved": "https://registry.npmjs.org/@opentelemetry/sdk-logs/-/sdk-logs-0.200.0.tgz", - "integrity": "sha512-VZG870063NLfObmQQNtCVcdXXLzI3vOjjrRENmU37HYiPFa0ZXpXVDsTD02Nh3AT3xYJzQaWKl2X2lQ2l7TWJA==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -11825,15 +11901,42 @@ "@opentelemetry/api": ">=1.4.0 <1.10.0" } }, - "node_modules/@opentelemetry/sdk-metrics": { + "node_modules/@opentelemetry/sdk-logs/node_modules/@opentelemetry/core": { + "version": "2.0.0", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@opentelemetry/semantic-conventions": "^1.29.0" + }, + "engines": { + "node": "^18.19.0 || >=20.6.0" + }, + "peerDependencies": { + "@opentelemetry/api": ">=1.0.0 <1.10.0" + } + }, + "node_modules/@opentelemetry/sdk-logs/node_modules/@opentelemetry/resources": { "version": "2.0.0", - "resolved": "https://registry.npmjs.org/@opentelemetry/sdk-metrics/-/sdk-metrics-2.0.0.tgz", - "integrity": "sha512-Bvy8QDjO05umd0+j+gDeWcTaVa1/R2lDj/eOvjzpm8VQj1K1vVZJuyjThpV5/lSHyYW2JaHF2IQ7Z8twJFAhjA==", "dev": true, "license": "Apache-2.0", "dependencies": { "@opentelemetry/core": "2.0.0", - "@opentelemetry/resources": "2.0.0" + "@opentelemetry/semantic-conventions": "^1.29.0" + }, + "engines": { + "node": "^18.19.0 || >=20.6.0" + }, + "peerDependencies": { + "@opentelemetry/api": ">=1.3.0 <1.10.0" + } + }, + "node_modules/@opentelemetry/sdk-metrics": { + "version": "2.0.1", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@opentelemetry/core": "2.0.1", + "@opentelemetry/resources": "2.0.1" }, "engines": { "node": "^18.19.0 || >=20.6.0" @@ -11844,8 +11947,6 @@ }, "node_modules/@opentelemetry/sdk-trace-base": { "version": "2.0.0", - "resolved": "https://registry.npmjs.org/@opentelemetry/sdk-trace-base/-/sdk-trace-base-2.0.0.tgz", - "integrity": "sha512-qQnYdX+ZCkonM7tA5iU4fSRsVxbFGml8jbxOgipRGMFHKaXKHQ30js03rTobYjKjIfnOsZSbHKWF0/0v0OQGfw==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -11860,10 +11961,37 @@ "@opentelemetry/api": ">=1.3.0 <1.10.0" } }, + "node_modules/@opentelemetry/sdk-trace-base/node_modules/@opentelemetry/core": { + "version": "2.0.0", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@opentelemetry/semantic-conventions": "^1.29.0" + }, + "engines": { + "node": "^18.19.0 || >=20.6.0" + }, + "peerDependencies": { + "@opentelemetry/api": ">=1.0.0 <1.10.0" + } + }, + "node_modules/@opentelemetry/sdk-trace-base/node_modules/@opentelemetry/resources": { + "version": "2.0.0", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@opentelemetry/core": "2.0.0", + "@opentelemetry/semantic-conventions": "^1.29.0" + }, + "engines": { + "node": "^18.19.0 || >=20.6.0" + }, + "peerDependencies": { + "@opentelemetry/api": ">=1.3.0 <1.10.0" + } + }, "node_modules/@opentelemetry/semantic-conventions": { "version": "1.33.0", - "resolved": "https://registry.npmjs.org/@opentelemetry/semantic-conventions/-/semantic-conventions-1.33.0.tgz", - "integrity": "sha512-TIpZvE8fiEILFfTlfPnltpBaD3d9/+uQHVCyC3vfdh6WfCXKhNFzoP5RyDDIndfvZC5GrA4pyEDNyjPloJud+w==", "dev": true, "license": "Apache-2.0", "engines": { @@ -11903,36 +12031,26 @@ }, "node_modules/@protobufjs/aspromise": { "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@protobufjs/aspromise/-/aspromise-1.1.2.tgz", - "integrity": "sha512-j+gKExEuLmKwvz3OgROXtrJ2UG2x8Ch2YZUxahh+s1F2HZ+wAceUNLkvy6zKCPVRkU++ZWQrdxsUeQXmcg4uoQ==", "dev": true, "license": "BSD-3-Clause" }, "node_modules/@protobufjs/base64": { "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@protobufjs/base64/-/base64-1.1.2.tgz", - "integrity": "sha512-AZkcAA5vnN/v4PDqKyMR5lx7hZttPDgClv83E//FMNhR2TMcLUhfRUBHCmSl0oi9zMgDDqRUJkSxO3wm85+XLg==", "dev": true, "license": "BSD-3-Clause" }, "node_modules/@protobufjs/codegen": { "version": "2.0.4", - "resolved": "https://registry.npmjs.org/@protobufjs/codegen/-/codegen-2.0.4.tgz", - "integrity": "sha512-YyFaikqM5sH0ziFZCN3xDC7zeGaB/d0IUb9CATugHWbd1FRFwWwt4ld4OYMPWu5a3Xe01mGAULCdqhMlPl29Jg==", "dev": true, "license": "BSD-3-Clause" }, "node_modules/@protobufjs/eventemitter": { "version": "1.1.0", - "resolved": "https://registry.npmjs.org/@protobufjs/eventemitter/-/eventemitter-1.1.0.tgz", - "integrity": "sha512-j9ednRT81vYJ9OfVuXG6ERSTdEL1xVsNgqpkxMsbIabzSo3goCjDIveeGv5d03om39ML71RdmrGNjG5SReBP/Q==", "dev": true, "license": "BSD-3-Clause" }, "node_modules/@protobufjs/fetch": { "version": "1.1.0", - "resolved": "https://registry.npmjs.org/@protobufjs/fetch/-/fetch-1.1.0.tgz", - "integrity": "sha512-lljVXpqXebpsijW71PZaCYeIcE5on1w5DlQy5WH6GLbFryLUrBD4932W/E2BSpfRJWseIL4v/KPgBFxDOIdKpQ==", "dev": true, "license": "BSD-3-Clause", "dependencies": { @@ -11942,36 +12060,26 @@ }, "node_modules/@protobufjs/float": { "version": "1.0.2", - "resolved": "https://registry.npmjs.org/@protobufjs/float/-/float-1.0.2.tgz", - "integrity": "sha512-Ddb+kVXlXst9d+R9PfTIxh1EdNkgoRe5tOX6t01f1lYWOvJnSPDBlG241QLzcyPdoNTsblLUdujGSE4RzrTZGQ==", "dev": true, "license": "BSD-3-Clause" }, "node_modules/@protobufjs/inquire": { "version": "1.1.0", - "resolved": "https://registry.npmjs.org/@protobufjs/inquire/-/inquire-1.1.0.tgz", - "integrity": "sha512-kdSefcPdruJiFMVSbn801t4vFK7KB/5gd2fYvrxhuJYg8ILrmn9SKSX2tZdV6V+ksulWqS7aXjBcRXl3wHoD9Q==", "dev": true, "license": "BSD-3-Clause" }, "node_modules/@protobufjs/path": { "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@protobufjs/path/-/path-1.1.2.tgz", - "integrity": "sha512-6JOcJ5Tm08dOHAbdR3GrvP+yUUfkjG5ePsHYczMFLq3ZmMkAD98cDgcT2iA1lJ9NVwFd4tH/iSSoe44YWkltEA==", "dev": true, "license": "BSD-3-Clause" }, "node_modules/@protobufjs/pool": { "version": "1.1.0", - "resolved": "https://registry.npmjs.org/@protobufjs/pool/-/pool-1.1.0.tgz", - "integrity": "sha512-0kELaGSIDBKvcgS4zkjz1PeddatrjYcmMWOlAuAPwAeccUrPHdUqo/J6LiymHHEiJT5NrF1UVwxY14f+fy4WQw==", "dev": true, "license": "BSD-3-Clause" }, "node_modules/@protobufjs/utf8": { "version": "1.1.0", - "resolved": "https://registry.npmjs.org/@protobufjs/utf8/-/utf8-1.1.0.tgz", - "integrity": "sha512-Vvn3zZrhQZkkBE8LSuW3em98c0FwgO4nxzv6OdSxPKJIEKY2bGbHn+mhGIPerzI4twdxaP8/0+06HBpwf345Lw==", "dev": true, "license": "BSD-3-Clause" }, @@ -19826,8 +19934,6 @@ }, "node_modules/long": { "version": "5.3.2", - "resolved": "https://registry.npmjs.org/long/-/long-5.3.2.tgz", - "integrity": "sha512-mNAgZ1GmyNhD7AuqnTG3/VQ26o760+ZYBPKjPvugO8+nLbYfX6TVpJPseBvopbdY+qpZ/lKUnmEc1LeZYS3QAA==", "dev": true, "license": "Apache-2.0" }, @@ -21557,9 +21663,7 @@ "license": "MIT" }, "node_modules/protobufjs": { - "version": "7.5.1", - "resolved": "https://registry.npmjs.org/protobufjs/-/protobufjs-7.5.1.tgz", - "integrity": "sha512-3qx3IRjR9WPQKagdwrKjO3Gu8RgQR2qqw+1KnigWhoVjFqegIj1K3bP11sGqhxrO46/XL7lekuG4jmjL+4cLsw==", + "version": "7.5.2", "dev": true, "hasInstallScript": true, "license": "BSD-3-Clause", @@ -23601,9 +23705,8 @@ }, "node_modules/ts-node": { "version": "10.9.2", - "resolved": "https://registry.npmjs.org/ts-node/-/ts-node-10.9.2.tgz", - "integrity": "sha512-f0FFpIdcHgn8zcPSbf1dRevwt047YMnaiJM3u2w2RewrB+fob/zePZcrOyQoLMMO7aBIddLcQIEK5dYjkLnGrQ==", "dev": true, + "license": "MIT", "dependencies": { "@cspotcode/source-map-support": "^0.8.0", "@tsconfig/node10": "^1.0.7", diff --git a/packages/amazonq/.changes/next-release/Bug Fix-2471c584-3904-4d90-bb7c-61efff219e43.json b/packages/amazonq/.changes/next-release/Bug Fix-2471c584-3904-4d90-bb7c-61efff219e43.json new file mode 100644 index 00000000000..f7af0fbb1a4 --- /dev/null +++ b/packages/amazonq/.changes/next-release/Bug Fix-2471c584-3904-4d90-bb7c-61efff219e43.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "Fix Error: 'Amazon Q service is not signed in'" +} diff --git a/packages/amazonq/.changes/next-release/Bug Fix-91d391d4-3777-4053-9e71-15b36dfa1f67.json b/packages/amazonq/.changes/next-release/Bug Fix-91d391d4-3777-4053-9e71-15b36dfa1f67.json new file mode 100644 index 00000000000..e3a608296a0 --- /dev/null +++ b/packages/amazonq/.changes/next-release/Bug Fix-91d391d4-3777-4053-9e71-15b36dfa1f67.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "Fix Error: 'Amazon Q Profile is not selected for IDC connection type'" +} diff --git a/packages/amazonq/.changes/next-release/Bug Fix-9b0e6490-39a8-445f-9d67-9d762de7421c.json b/packages/amazonq/.changes/next-release/Bug Fix-9b0e6490-39a8-445f-9d67-9d762de7421c.json new file mode 100644 index 00000000000..f17516bb8f4 --- /dev/null +++ b/packages/amazonq/.changes/next-release/Bug Fix-9b0e6490-39a8-445f-9d67-9d762de7421c.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "Previous and subsequent cells are used as context for completion in a Jupyter notebook" +} diff --git a/packages/amazonq/.changes/next-release/Bug Fix-bb976b5f-7697-42d8-89a9-8e96310a23f4.json b/packages/amazonq/.changes/next-release/Bug Fix-bb976b5f-7697-42d8-89a9-8e96310a23f4.json new file mode 100644 index 00000000000..988fb2bcc7b --- /dev/null +++ b/packages/amazonq/.changes/next-release/Bug Fix-bb976b5f-7697-42d8-89a9-8e96310a23f4.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "Support chat in AL2 aarch64" +} diff --git a/packages/amazonq/.changes/next-release/Feature-31d91f84-30cb-4acd-9e39-9dc153edf0a6.json b/packages/amazonq/.changes/next-release/Feature-31d91f84-30cb-4acd-9e39-9dc153edf0a6.json new file mode 100644 index 00000000000..da0d200410d --- /dev/null +++ b/packages/amazonq/.changes/next-release/Feature-31d91f84-30cb-4acd-9e39-9dc153edf0a6.json @@ -0,0 +1,4 @@ +{ + "type": "Feature", + "description": "Add inline completion support for abap language" +} diff --git a/packages/amazonq/.changes/next-release/bugfix-ad9a8829-efa0-463b-ba2f-c89ff1cb69e4.json b/packages/amazonq/.changes/next-release/bugfix-ad9a8829-efa0-463b-ba2f-c89ff1cb69e4.json new file mode 100644 index 00000000000..1b84a3f57c0 --- /dev/null +++ b/packages/amazonq/.changes/next-release/bugfix-ad9a8829-efa0-463b-ba2f-c89ff1cb69e4.json @@ -0,0 +1,4 @@ +{ + "type": "bugfix", + "description": "/review: disable auto-review by default" +} diff --git a/packages/amazonq/package.json b/packages/amazonq/package.json index 7a3aab1f8d6..5820ecdff8c 100644 --- a/packages/amazonq/package.json +++ b/packages/amazonq/package.json @@ -59,7 +59,6 @@ "watch": "npm run clean && npm run buildScripts && tsc -watch -p ./", "testCompile": "npm run clean && npm run buildScripts && npm run compileOnly", "test": "npm run testCompile && c8 --allowExternal ts-node ../core/scripts/test/launchTest.ts unit dist/test/unit/index.js ../core/dist/src/testFixtures/workspaceFolder", - "testE2E": "npm run testCompile && c8 --allowExternal ts-node ../core/scripts/test/launchTest.ts e2e dist/test/e2e/index.js ../core/dist/src/testFixtures/workspaceFolder", "testWeb": "npm run compileDev && c8 --allowExternal ts-node ../core/scripts/test/launchTest.ts web dist/test/web/testRunnerWebCore.js", "webRun": "npx @vscode/test-web --open-devtools --browserOption=--disable-web-security --waitForDebugger=9222 --extensionDevelopmentPath=. .", "webWatch": "npm run clean && npm run buildScripts && webpack --mode development --watch", diff --git a/packages/amazonq/src/api.ts b/packages/amazonq/src/api.ts index 03b2a32ea55..bd7d5c6a361 100644 --- a/packages/amazonq/src/api.ts +++ b/packages/amazonq/src/api.ts @@ -8,6 +8,7 @@ import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseReques import { AuthUtil } from 'aws-core-vscode/codewhisperer' import { ChatSession } from 'aws-core-vscode/codewhispererChat' import { api } from 'aws-core-vscode/amazonq' +import { getLogger } from 'aws-core-vscode/shared' export default { chatApi: { @@ -26,8 +27,25 @@ export default { await AuthUtil.instance.showReauthenticatePrompt() } }, + /** + * @deprecated use getAuthState() instead + * + * Legacy function for callers who expect auth state to be granular amongst Q features. + * Auth state is consistent between features, so getAuthState() can be consumed safely for all features. + * + */ async getChatAuthState() { - return AuthUtil.instance.getChatAuthState() + getLogger().warn('Warning: getChatAuthState() is deprecated. Use getAuthState() instead.') + const state = AuthUtil.instance.getAuthState() + const convertedState = state === 'notConnected' ? 'disconnected' : state + return { + codewhispererCore: convertedState, + codewhispererChat: convertedState, + amazonQ: convertedState, + } + }, + getAuthState() { + return AuthUtil.instance.getAuthState() }, }, } satisfies api diff --git a/packages/amazonq/src/app/amazonqScan/app.ts b/packages/amazonq/src/app/amazonqScan/app.ts index 21857163bd2..66f45246129 100644 --- a/packages/amazonq/src/app/amazonqScan/app.ts +++ b/packages/amazonq/src/app/amazonqScan/app.ts @@ -19,6 +19,7 @@ import { Messenger } from './chat/controller/messenger/messenger' import { UIMessageListener } from './chat/views/actions/uiMessageListener' import { debounce } from 'lodash' import { Commands, placeholder } from 'aws-core-vscode/shared' +import { auth2 } from 'aws-core-vscode/auth' export function init(appContext: AmazonQAppInitContext) { const scanChatControllerEventEmitters: ScanChatControllerEventEmitters = { @@ -52,7 +53,7 @@ export function init(appContext: AmazonQAppInitContext) { appContext.registerWebViewToAppMessagePublisher(new MessagePublisher(scanChatUIInputEventEmitter), 'review') const debouncedEvent = debounce(async () => { - const authenticated = (await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected' + const authenticated = AuthUtil.instance.getAuthState() === 'connected' let authenticatingSessionID = '' if (authenticated) { @@ -67,7 +68,7 @@ export function init(appContext: AmazonQAppInitContext) { messenger.sendAuthenticationUpdate(authenticated, [authenticatingSessionID]) }, 500) - AuthUtil.instance.secondaryAuth.onDidChangeActiveConnection(() => { + AuthUtil.instance.onDidChangeConnectionState((e: auth2.AuthStateEvent) => { return debouncedEvent() }) AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(() => { diff --git a/packages/amazonq/src/app/amazonqScan/chat/controller/controller.ts b/packages/amazonq/src/app/amazonqScan/chat/controller/controller.ts index 72af0a200c5..1262bce32ee 100644 --- a/packages/amazonq/src/app/amazonqScan/chat/controller/controller.ts +++ b/packages/amazonq/src/app/amazonqScan/chat/controller/controller.ts @@ -104,7 +104,7 @@ export class ScanController { telemetry.amazonq_feedback.emit({ featureId: 'amazonQReview', amazonqConversationId: this.sessionStorage.getSession().scanUuid, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, interactionType: data.vote, }) }) @@ -122,8 +122,8 @@ export class ScanController { try { getLogger().debug(`Q - Review: Session created with id: ${session.tabID}`) - const authState = await AuthUtil.instance.getChatAuthState() - if (authState.amazonQ !== 'connected') { + const authState = AuthUtil.instance.getAuthState() + if (authState !== 'connected') { void this.messenger.sendAuthNeededExceptionMessage(authState, tabID) session.isAuthenticating = true return @@ -161,8 +161,8 @@ export class ScanController { return } // check that the session is authenticated - const authState = await AuthUtil.instance.getChatAuthState() - if (authState.amazonQ !== 'connected') { + const authState = AuthUtil.instance.getAuthState() + if (authState !== 'connected') { void this.messenger.sendAuthNeededExceptionMessage(authState, message.tabID) session.isAuthenticating = true return diff --git a/packages/amazonq/src/app/amazonqScan/chat/controller/messenger/messenger.ts b/packages/amazonq/src/app/amazonqScan/chat/controller/messenger/messenger.ts index 18b05e8bb84..d165fb0b46c 100644 --- a/packages/amazonq/src/app/amazonqScan/chat/controller/messenger/messenger.ts +++ b/packages/amazonq/src/app/amazonqScan/chat/controller/messenger/messenger.ts @@ -10,7 +10,6 @@ import { AuthFollowUpType, AuthMessageDataMap } from 'aws-core-vscode/amazonq' import { - FeatureAuthState, SecurityScanError, CodeWhispererConstants, SecurityScanStep, @@ -34,6 +33,7 @@ import { import { i18n } from 'aws-core-vscode/shared' import { ScanAction, scanProgressMessage } from '../../../models/constants' import path from 'path' +import { auth2 } from 'aws-core-vscode/auth' export type UnrecoverableErrorType = 'no-project-found' | 'no-open-file-found' | 'invalid-file-type' @@ -78,19 +78,15 @@ export class Messenger { this.dispatcher.sendUpdatePromptProgress(new UpdatePromptProgressMessage(tabID, progressField)) } - public async sendAuthNeededExceptionMessage(credentialState: FeatureAuthState, tabID: string) { + public async sendAuthNeededExceptionMessage(credentialState: auth2.AuthState, tabID: string) { let authType: AuthFollowUpType = 'full-auth' let message = AuthMessageDataMap[authType].message - switch (credentialState.amazonQ) { - case 'disconnected': + switch (credentialState) { + case 'notConnected': authType = 'full-auth' message = AuthMessageDataMap[authType].message break - case 'unsupported': - authType = 'use-supported-auth' - message = AuthMessageDataMap[authType].message - break case 'expired': authType = 're-auth' message = AuthMessageDataMap[authType].message diff --git a/packages/amazonq/src/app/chat/activation.ts b/packages/amazonq/src/app/chat/activation.ts index bf6b7cdc3df..af48bc65e05 100644 --- a/packages/amazonq/src/app/chat/activation.ts +++ b/packages/amazonq/src/app/chat/activation.ts @@ -6,22 +6,14 @@ import * as vscode from 'vscode' import { ExtensionContext } from 'vscode' import { telemetry } from 'aws-core-vscode/telemetry' -import { AuthUtil, CodeWhispererSettings } from 'aws-core-vscode/codewhisperer' -import { Commands, placeholder, funcUtil } from 'aws-core-vscode/shared' +import { AuthUtil } from 'aws-core-vscode/codewhisperer' +import { Commands, placeholder } from 'aws-core-vscode/shared' import * as amazonq from 'aws-core-vscode/amazonq' export async function activate(context: ExtensionContext) { const appInitContext = amazonq.DefaultAmazonQAppInitContext.instance await amazonq.TryChatCodeLensProvider.register(appInitContext.onDidChangeAmazonQVisibility.event) - const setupLsp = funcUtil.debounce(async () => { - void amazonq.LspController.instance.trySetupLsp(context, { - startUrl: AuthUtil.instance.startUrl, - maxIndexSize: CodeWhispererSettings.instance.getMaxIndexSize(), - isVectorIndexEnabled: false, - }) - }, 5000) - context.subscriptions.push( amazonq.focusAmazonQChatWalkthrough.register(), amazonq.walkthroughInlineSuggestionsExample.register(), @@ -37,7 +29,6 @@ export async function activate(context: ExtensionContext) { void vscode.env.openExternal(vscode.Uri.parse(amazonq.amazonQHelpUrl)) }) - void setupLsp() void setupAuthNotification() } diff --git a/packages/amazonq/src/app/inline/activation.ts b/packages/amazonq/src/app/inline/activation.ts index d786047b2aa..69515127441 100644 --- a/packages/amazonq/src/app/inline/activation.ts +++ b/packages/amazonq/src/app/inline/activation.ts @@ -6,68 +6,26 @@ import vscode from 'vscode' import { AuthUtil, - CodeSuggestionsState, - CodeWhispererCodeCoverageTracker, CodeWhispererConstants, - CodeWhispererSettings, - ConfigurationEntry, - DefaultCodeWhispererClient, - invokeRecommendation, isInlineCompletionEnabled, - KeyStrokeHandler, - RecommendationHandler, runtimeLanguageContext, TelemetryHelper, UserWrittenCodeTracker, vsCodeState, } from 'aws-core-vscode/codewhisperer' -import { Commands, getLogger, globals, sleep } from 'aws-core-vscode/shared' +import { globals, sleep } from 'aws-core-vscode/shared' export async function activate() { - const codewhispererSettings = CodeWhispererSettings.instance - const client = new DefaultCodeWhispererClient() - if (isInlineCompletionEnabled()) { await setSubscriptionsforInlineCompletion() await AuthUtil.instance.setVscodeContextProps() } - function getAutoTriggerStatus(): boolean { - return CodeSuggestionsState.instance.isSuggestionsEnabled() - } - - async function getConfigEntry(): Promise { - const isShowMethodsEnabled: boolean = - vscode.workspace.getConfiguration('editor').get('suggest.showMethods') || false - const isAutomatedTriggerEnabled: boolean = getAutoTriggerStatus() - const isManualTriggerEnabled: boolean = true - const isSuggestionsWithCodeReferencesEnabled = codewhispererSettings.isSuggestionsWithCodeReferencesEnabled() - - // TODO:remove isManualTriggerEnabled - return { - isShowMethodsEnabled, - isManualTriggerEnabled, - isAutomatedTriggerEnabled, - isSuggestionsWithCodeReferencesEnabled, - } - } - async function setSubscriptionsforInlineCompletion() { - RecommendationHandler.instance.subscribeSuggestionCommands() - /** * Automated trigger */ globals.context.subscriptions.push( - vscode.window.onDidChangeActiveTextEditor(async (editor) => { - await RecommendationHandler.instance.onEditorChange() - }), - vscode.window.onDidChangeWindowState(async (e) => { - await RecommendationHandler.instance.onFocusChange() - }), - vscode.window.onDidChangeTextEditorSelection(async (e) => { - await RecommendationHandler.instance.onCursorChange(e) - }), vscode.workspace.onDidChangeTextDocument(async (e) => { const editor = vscode.window.activeTextEditor if (!editor) { @@ -80,7 +38,6 @@ export async function activate() { return } - CodeWhispererCodeCoverageTracker.getTracker(e.document.languageId)?.countTotalTokens(e) UserWrittenCodeTracker.instance.onTextDocumentChange(e) /** * Handle this keystroke event only when @@ -105,19 +62,6 @@ export async function activate() { * Then this event can be processed by our code. */ await sleep(CodeWhispererConstants.vsCodeCursorUpdateDelay) - if (!RecommendationHandler.instance.isSuggestionVisible()) { - await KeyStrokeHandler.instance.processKeyStroke(e, editor, client, await getConfigEntry()) - } - }), - // manual trigger - Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => { - invokeRecommendation( - vscode.window.activeTextEditor as vscode.TextEditor, - client, - await getConfigEntry() - ).catch((e) => { - getLogger().error('invokeRecommendation failed: %s', (e as Error).message) - }) }) ) } diff --git a/packages/amazonq/src/app/inline/completion.ts b/packages/amazonq/src/app/inline/completion.ts index be390cef34c..78e0fa9cd13 100644 --- a/packages/amazonq/src/app/inline/completion.ts +++ b/packages/amazonq/src/app/inline/completion.ts @@ -8,7 +8,6 @@ import { InlineCompletionContext, InlineCompletionItem, InlineCompletionItemProvider, - InlineCompletionList, Position, TextDocument, commands, @@ -16,6 +15,8 @@ import { Disposable, window, TextEditor, + InlineCompletionTriggerKind, + Range, } from 'vscode' import { LanguageClient } from 'vscode-languageclient' import { @@ -27,10 +28,20 @@ import { RecommendationService } from './recommendationService' import { CodeWhispererConstants, ReferenceHoverProvider, - ReferenceInlineProvider, ReferenceLogViewProvider, ImportAdderProvider, + CodeSuggestionsState, + vsCodeState, + inlineCompletionsDebounceDelay, + noInlineSuggestionsMsg, + ReferenceInlineProvider, } from 'aws-core-vscode/codewhisperer' +import { InlineGeneratingMessage } from './inlineGeneratingMessage' +import { LineTracker } from './stateTracker/lineTracker' +import { InlineTutorialAnnotation } from './tutorials/inlineTutorialAnnotation' +import { TelemetryHelper } from './telemetryHelper' +import { getLogger } from 'aws-core-vscode/shared' +import { debounce, messageUtils } from 'aws-core-vscode/utils' export class InlineCompletionManager implements Disposable { private disposable: Disposable @@ -38,26 +49,42 @@ export class InlineCompletionManager implements Disposable { private languageClient: LanguageClient private sessionManager: SessionManager private recommendationService: RecommendationService + private lineTracker: LineTracker + private incomingGeneratingMessage: InlineGeneratingMessage + private inlineTutorialAnnotation: InlineTutorialAnnotation private readonly logSessionResultMessageName = 'aws/logInlineCompletionSessionResults' - constructor(languageClient: LanguageClient) { + constructor( + languageClient: LanguageClient, + sessionManager: SessionManager, + lineTracker: LineTracker, + inlineTutorialAnnotation: InlineTutorialAnnotation + ) { this.languageClient = languageClient - this.sessionManager = new SessionManager() - this.recommendationService = new RecommendationService(this.sessionManager) + this.sessionManager = sessionManager + this.lineTracker = lineTracker + this.incomingGeneratingMessage = new InlineGeneratingMessage(this.lineTracker) + this.recommendationService = new RecommendationService(this.sessionManager, this.incomingGeneratingMessage) + this.inlineTutorialAnnotation = inlineTutorialAnnotation this.inlineCompletionProvider = new AmazonQInlineCompletionItemProvider( languageClient, this.recommendationService, - this.sessionManager + this.sessionManager, + this.inlineTutorialAnnotation ) this.disposable = languages.registerInlineCompletionItemProvider( CodeWhispererConstants.platformLanguageIds, this.inlineCompletionProvider ) + + this.lineTracker.ready() } public dispose(): void { if (this.disposable) { this.disposable.dispose() + this.incomingGeneratingMessage.dispose() + this.lineTracker.dispose() } } @@ -97,10 +124,21 @@ export class InlineCompletionManager implements Disposable { ) ReferenceLogViewProvider.instance.addReferenceLog(referenceLog) ReferenceHoverProvider.instance.addCodeReferences(item.insertText as string, item.references) + + // Show codelense for 5 seconds. + ReferenceInlineProvider.instance.setInlineReference( + startLine, + item.insertText as string, + item.references + ) + setTimeout(() => { + ReferenceInlineProvider.instance.removeInlineReference() + }, 5000) } if (item.mostRelevantMissingImports?.length) { await ImportAdderProvider.instance.onAcceptRecommendation(editor, item, startLine) } + this.sessionManager.incrementSuggestionCount() } commands.registerCommand('aws.amazonq.acceptInline', onInlineAcceptance) @@ -130,38 +168,6 @@ export class InlineCompletionManager implements Disposable { this.languageClient.sendNotification(this.logSessionResultMessageName, params) } commands.registerCommand('aws.amazonq.rejectCodeSuggestion', onInlineRejection) - - /* - We have to overwrite the prev. and next. commands because the inlineCompletionProvider only contained the current item - To show prev. and next. recommendation we need to re-register a new provider with the previous or next item - */ - - const swapProviderAndShow = async () => { - await commands.executeCommand('editor.action.inlineSuggest.hide') - this.disposable.dispose() - this.disposable = languages.registerInlineCompletionItemProvider( - CodeWhispererConstants.platformLanguageIds, - new AmazonQInlineCompletionItemProvider( - this.languageClient, - this.recommendationService, - this.sessionManager, - false - ) - ) - await commands.executeCommand('editor.action.inlineSuggest.trigger') - } - - const prevCommandHandler = async () => { - this.sessionManager.decrementActiveIndex() - await swapProviderAndShow() - } - commands.registerCommand('editor.action.inlineSuggest.showPrevious', prevCommandHandler) - - const nextCommandHandler = async () => { - this.sessionManager.incrementActiveIndex() - await swapProviderAndShow() - } - commands.registerCommand('editor.action.inlineSuggest.showNext', nextCommandHandler) } } @@ -170,17 +176,34 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem private readonly languageClient: LanguageClient, private readonly recommendationService: RecommendationService, private readonly sessionManager: SessionManager, - private readonly isNewSession: boolean = true + private readonly inlineTutorialAnnotation: InlineTutorialAnnotation ) {} - async provideInlineCompletionItems( + provideInlineCompletionItems = debounce( + this._provideInlineCompletionItems.bind(this), + inlineCompletionsDebounceDelay, + true + ) + + private async _provideInlineCompletionItems( document: TextDocument, position: Position, context: InlineCompletionContext, token: CancellationToken - ): Promise { - if (this.isNewSession) { - // make service requests if it's a new session + ): Promise { + try { + vsCodeState.isRecommendationsActive = true + const isAutoTrigger = context.triggerKind === InlineCompletionTriggerKind.Automatic + if (isAutoTrigger && !CodeSuggestionsState.instance.isSuggestionsEnabled()) { + // return early when suggestions are disabled with auto trigger + return [] + } + + // tell the tutorial that completions has been triggered + await this.inlineTutorialAnnotation.triggered(context.triggerKind) + TelemetryHelper.instance.setInvokeSuggestionStartTime() + TelemetryHelper.instance.setTriggerType(context.triggerKind) + await this.recommendationService.getAllRecommendations( this.languageClient, document, @@ -188,34 +211,47 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem context, token ) - } - // get active item from session for displaying - const items = this.sessionManager.getActiveRecommendation() - const session = this.sessionManager.getActiveSession() - if (!session || !items.length) { - return [] - } - const editor = window.activeTextEditor - for (const item of items) { - item.command = { - command: 'aws.amazonq.acceptInline', - title: 'On acceptance', - arguments: [ - session.sessionId, - item, - editor, - session.requestStartTime, - position.line, - session.firstCompletionDisplayLatency, - ], + // get active item from session for displaying + const items = this.sessionManager.getActiveRecommendation() + const session = this.sessionManager.getActiveSession() + const editor = window.activeTextEditor + + // Show message to user when manual invoke fails to produce results. + if (items.length === 0 && context.triggerKind === InlineCompletionTriggerKind.Invoke) { + void messageUtils.showTimedMessage(noInlineSuggestionsMsg, 2000) } - ReferenceInlineProvider.instance.setInlineReference( - position.line, - item.insertText as string, - item.references - ) - ImportAdderProvider.instance.onShowRecommendation(document, position.line, item) + + if (!session || !items.length || !editor) { + getLogger().debug( + `Failed to produce inline suggestion results. Received ${items.length} items from service` + ) + return [] + } + + const cursorPosition = document.validatePosition(position) + for (const item of items) { + item.command = { + command: 'aws.amazonq.acceptInline', + title: 'On acceptance', + arguments: [ + session.sessionId, + item, + editor, + session.requestStartTime, + cursorPosition.line, + session.firstCompletionDisplayLatency, + ], + } + item.range = new Range(cursorPosition, cursorPosition) + item.insertText = typeof item.insertText === 'string' ? item.insertText : item.insertText.value + ImportAdderProvider.instance.onShowRecommendation(document, cursorPosition.line, item) + } + return items as InlineCompletionItem[] + } catch (e) { + getLogger('amazonqLsp').error('Failed to provide completion items: %O', e) + return [] + } finally { + vsCodeState.isRecommendationsActive = false } - return items as InlineCompletionItem[] } } diff --git a/packages/amazonq/src/app/inline/inlineGeneratingMessage.ts b/packages/amazonq/src/app/inline/inlineGeneratingMessage.ts new file mode 100644 index 00000000000..a9fe7a1d98c --- /dev/null +++ b/packages/amazonq/src/app/inline/inlineGeneratingMessage.ts @@ -0,0 +1,93 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import { editorUtilities } from 'aws-core-vscode/shared' +import * as vscode from 'vscode' +import { LineSelection, LineTracker } from './stateTracker/lineTracker' +import { AuthUtil } from 'aws-core-vscode/codewhisperer' +import { cancellableDebounce } from 'aws-core-vscode/utils' + +/** + * Manages the inline ghost text message show when Inline Suggestions is "thinking". + */ +export class InlineGeneratingMessage implements vscode.Disposable { + private readonly _disposable: vscode.Disposable + + private readonly cwLineHintDecoration: vscode.TextEditorDecorationType = + vscode.window.createTextEditorDecorationType({ + after: { + margin: '0 0 0 3em', + contentText: 'Amazon Q is generating...', + textDecoration: 'none', + fontWeight: 'normal', + fontStyle: 'normal', + color: 'var(--vscode-editorCodeLens-foreground)', + }, + rangeBehavior: vscode.DecorationRangeBehavior.OpenOpen, + isWholeLine: true, + }) + + constructor(private readonly lineTracker: LineTracker) { + this._disposable = vscode.Disposable.from( + AuthUtil.instance.onDidChangeConnectionState(async (e) => { + this.hideGenerating() + }) + ) + } + + dispose() { + this._disposable.dispose() + } + + readonly refreshDebounced = cancellableDebounce(async () => { + await this._refresh(true) + }, 1000) + + async showGenerating(triggerType: vscode.InlineCompletionTriggerKind) { + if (triggerType === vscode.InlineCompletionTriggerKind.Invoke) { + // if user triggers on demand, immediately update the UI and cancel the previous debounced update if there is one + this.refreshDebounced.cancel() + await this._refresh(true) + } else { + await this.refreshDebounced.promise() + } + } + + async _refresh(shouldDisplay: boolean) { + const editor = vscode.window.activeTextEditor + if (!editor) { + return + } + + const selections = this.lineTracker.selections + if (!editor || !selections || !editorUtilities.isTextEditor(editor)) { + this.hideGenerating() + return + } + + if (!AuthUtil.instance.isConnected()) { + this.hideGenerating() + return + } + + await this.updateDecorations(editor, selections, shouldDisplay) + } + + hideGenerating() { + vscode.window.activeTextEditor?.setDecorations(this.cwLineHintDecoration, []) + } + + async updateDecorations(editor: vscode.TextEditor, lines: LineSelection[], shouldDisplay: boolean) { + const range = editor.document.validateRange( + new vscode.Range(lines[0].active, lines[0].active, lines[0].active, lines[0].active) + ) + + if (shouldDisplay) { + editor.setDecorations(this.cwLineHintDecoration, [range]) + } else { + editor.setDecorations(this.cwLineHintDecoration, []) + } + } +} diff --git a/packages/amazonq/src/app/inline/recommendationService.ts b/packages/amazonq/src/app/inline/recommendationService.ts index 45dd0099ebd..ac0b16f140c 100644 --- a/packages/amazonq/src/app/inline/recommendationService.ts +++ b/packages/amazonq/src/app/inline/recommendationService.ts @@ -11,9 +11,15 @@ import { import { CancellationToken, InlineCompletionContext, Position, TextDocument } from 'vscode' import { LanguageClient } from 'vscode-languageclient' import { SessionManager } from './sessionManager' +import { InlineGeneratingMessage } from './inlineGeneratingMessage' +import { CodeWhispererStatusBarManager } from 'aws-core-vscode/codewhisperer' +import { TelemetryHelper } from './telemetryHelper' export class RecommendationService { - constructor(private readonly sessionManager: SessionManager) {} + constructor( + private readonly sessionManager: SessionManager, + private readonly inlineGeneratingMessage: InlineGeneratingMessage + ) {} async getAllRecommendations( languageClient: LanguageClient, @@ -30,29 +36,55 @@ export class RecommendationService { context, } const requestStartTime = Date.now() + const statusBar = CodeWhispererStatusBarManager.instance + TelemetryHelper.instance.setInvokeSuggestionStartTime() + TelemetryHelper.instance.setPreprocessEndTime() + TelemetryHelper.instance.setSdkApiCallStartTime() - // Handle first request - const firstResult: InlineCompletionListWithReferences = await languageClient.sendRequest( - inlineCompletionWithReferencesRequestType as any, - request, - token - ) + try { + // Show UI indicators that we are generating suggestions + await this.inlineGeneratingMessage.showGenerating(context.triggerKind) + await statusBar.setLoading() - const firstCompletionDisplayLatency = Date.now() - requestStartTime - this.sessionManager.startSession( - firstResult.sessionId, - firstResult.items, - requestStartTime, - firstCompletionDisplayLatency - ) + // Handle first request + const firstResult: InlineCompletionListWithReferences = await languageClient.sendRequest( + inlineCompletionWithReferencesRequestType.method, + request, + token + ) - if (firstResult.partialResultToken) { - // If there are more results to fetch, handle them in the background - this.processRemainingRequests(languageClient, request, firstResult, token).catch((error) => { - languageClient.warn(`Error when getting suggestions: ${error}`) - }) - } else { - this.sessionManager.closeSession() + // Set telemetry data for the first response + TelemetryHelper.instance.setSdkApiCallEndTime() + TelemetryHelper.instance.setSessionId(firstResult.sessionId) + if (firstResult.items.length > 0) { + TelemetryHelper.instance.setFirstResponseRequestId(firstResult.items[0].itemId) + } + TelemetryHelper.instance.setFirstSuggestionShowTime() + + const firstCompletionDisplayLatency = Date.now() - requestStartTime + this.sessionManager.startSession( + firstResult.sessionId, + firstResult.items, + requestStartTime, + firstCompletionDisplayLatency + ) + + if (firstResult.partialResultToken) { + // If there are more results to fetch, handle them in the background + this.processRemainingRequests(languageClient, request, firstResult, token).catch((error) => { + languageClient.warn(`Error when getting suggestions: ${error}`) + }) + } else { + this.sessionManager.closeSession() + + // No more results to fetch, mark pagination as complete + TelemetryHelper.instance.setAllPaginationEndTime() + TelemetryHelper.instance.tryRecordClientComponentLatency() + } + } finally { + // Remove all UI indicators of message generation since we are done + this.inlineGeneratingMessage.hideGenerating() + void statusBar.refreshStatusBar() // effectively "stop loading" } } @@ -66,13 +98,18 @@ export class RecommendationService { while (nextToken) { const request = { ...initialRequest, partialResultToken: nextToken } const result: InlineCompletionListWithReferences = await languageClient.sendRequest( - inlineCompletionWithReferencesRequestType as any, + inlineCompletionWithReferencesRequestType.method, request, token ) this.sessionManager.updateSessionSuggestions(result.items) nextToken = result.partialResultToken } + this.sessionManager.closeSession() + + // All pagination requests completed + TelemetryHelper.instance.setAllPaginationEndTime() + TelemetryHelper.instance.tryRecordClientComponentLatency() } } diff --git a/packages/amazonq/src/app/inline/sessionManager.ts b/packages/amazonq/src/app/inline/sessionManager.ts index 4b70a684001..10f436344bc 100644 --- a/packages/amazonq/src/app/inline/sessionManager.ts +++ b/packages/amazonq/src/app/inline/sessionManager.ts @@ -17,7 +17,8 @@ interface CodeWhispererSession { export class SessionManager { private activeSession?: CodeWhispererSession - private activeIndex: number = 0 + private _acceptedSuggestionCount: number = 0 + constructor() {} public startSession( @@ -33,7 +34,6 @@ export class SessionManager { requestStartTime, firstCompletionDisplayLatency, } - this.activeIndex = 0 } public closeSession() { @@ -54,49 +54,19 @@ export class SessionManager { this.activeSession.suggestions = [...this.activeSession.suggestions, ...suggestions] } - public incrementActiveIndex() { - const suggestionCount = this.activeSession?.suggestions?.length - if (!suggestionCount) { - return - } - this.activeIndex === suggestionCount - 1 ? suggestionCount - 1 : this.activeIndex++ + public getActiveRecommendation(): InlineCompletionItemWithReferences[] { + return this.activeSession?.suggestions ?? [] } - public decrementActiveIndex() { - this.activeIndex === 0 ? 0 : this.activeIndex-- + public get acceptedSuggestionCount(): number { + return this._acceptedSuggestionCount } - /* - We have to maintain the active suggestion index ourselves because VS Code doesn't expose which suggestion it's currently showing - In order to keep track of the right suggestion state, and for features such as reference tracker, this hack is still needed - */ - - public getActiveRecommendation(): InlineCompletionItemWithReferences[] { - let suggestionCount = this.activeSession?.suggestions.length - if (!suggestionCount) { - return [] - } - if (suggestionCount === 1 && this.activeSession?.isRequestInProgress) { - suggestionCount += 1 - } - - const activeSuggestion = this.activeSession?.suggestions[this.activeIndex] - if (!activeSuggestion) { - return [] - } - const items = [activeSuggestion] - // to make the total number of suggestions match the actual number - for (let i = 1; i < suggestionCount; i++) { - items.push({ - ...activeSuggestion, - insertText: `${i}`, - }) - } - return items + public incrementSuggestionCount() { + this._acceptedSuggestionCount += 1 } public clear() { this.activeSession = undefined - this.activeIndex = 0 } } diff --git a/packages/amazonq/src/app/inline/stateTracker/lineTracker.ts b/packages/amazonq/src/app/inline/stateTracker/lineTracker.ts new file mode 100644 index 00000000000..58bee329a40 --- /dev/null +++ b/packages/amazonq/src/app/inline/stateTracker/lineTracker.ts @@ -0,0 +1,178 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as vscode from 'vscode' +import { editorUtilities, setContext } from 'aws-core-vscode/shared' + +export interface LineSelection { + anchor: number + active: number +} + +export interface LinesChangeEvent { + readonly editor: vscode.TextEditor | undefined + readonly selections: LineSelection[] | undefined + + readonly reason: 'editor' | 'selection' | 'content' +} + +/** + * This class providees a single interface to manage and access users' "line" selections + * Callers could use it by subscribing onDidChangeActiveLines to do UI updates or logic needed to be executed when line selections get changed + */ +export class LineTracker implements vscode.Disposable { + private _onDidChangeActiveLines = new vscode.EventEmitter() + get onDidChangeActiveLines(): vscode.Event { + return this._onDidChangeActiveLines.event + } + + private _editor: vscode.TextEditor | undefined + private _disposable: vscode.Disposable | undefined + + private _selections: LineSelection[] | undefined + get selections(): LineSelection[] | undefined { + return this._selections + } + + private _onReady: vscode.EventEmitter = new vscode.EventEmitter() + get onReady(): vscode.Event { + return this._onReady.event + } + + private _ready: boolean = false + get isReady() { + return this._ready + } + + constructor() { + this._disposable = vscode.Disposable.from( + vscode.window.onDidChangeActiveTextEditor(async (e) => { + await this.onActiveTextEditorChanged(e) + }), + vscode.window.onDidChangeTextEditorSelection(async (e) => { + await this.onTextEditorSelectionChanged(e) + }), + vscode.workspace.onDidChangeTextDocument((e) => { + this.onContentChanged(e) + }) + ) + + queueMicrotask(async () => await this.onActiveTextEditorChanged(vscode.window.activeTextEditor)) + } + + dispose() { + this._disposable?.dispose() + } + + ready() { + if (this._ready) { + throw new Error('Linetracker is already activated') + } + + this._ready = true + queueMicrotask(() => this._onReady.fire()) + } + + // @VisibleForTesting + async onActiveTextEditorChanged(editor: vscode.TextEditor | undefined) { + if (editor === this._editor) { + return + } + + this._editor = editor + this._selections = toLineSelections(editor?.selections) + if (this._selections && this._selections[0]) { + const s = this._selections.map((item) => item.active + 1) + await setContext('codewhisperer.activeLine', s) + } + + this.notifyLinesChanged('editor') + } + + // @VisibleForTesting + async onTextEditorSelectionChanged(e: vscode.TextEditorSelectionChangeEvent) { + // If this isn't for our cached editor and its not a real editor -- kick out + if (this._editor !== e.textEditor && !editorUtilities.isTextEditor(e.textEditor)) { + return + } + + const selections = toLineSelections(e.selections) + if (this._editor === e.textEditor && this.includes(selections)) { + return + } + + this._editor = e.textEditor + this._selections = selections + if (this._selections && this._selections[0]) { + const s = this._selections.map((item) => item.active + 1) + await setContext('codewhisperer.activeLine', s) + } + + this.notifyLinesChanged('selection') + } + + // @VisibleForTesting + onContentChanged(e: vscode.TextDocumentChangeEvent) { + const editor = vscode.window.activeTextEditor + if (e.document === editor?.document && e.contentChanges.length > 0 && editorUtilities.isTextEditor(editor)) { + this._editor = editor + this._selections = toLineSelections(this._editor?.selections) + + this.notifyLinesChanged('content') + } + } + + notifyLinesChanged(reason: 'editor' | 'selection' | 'content') { + const e: LinesChangeEvent = { editor: this._editor, selections: this.selections, reason: reason } + this._onDidChangeActiveLines.fire(e) + } + + includes(selections: LineSelection[]): boolean + includes(line: number, options?: { activeOnly: boolean }): boolean + includes(lineOrSelections: number | LineSelection[], options?: { activeOnly: boolean }): boolean { + if (typeof lineOrSelections !== 'number') { + return isIncluded(lineOrSelections, this._selections) + } + + if (this._selections === undefined || this._selections.length === 0) { + return false + } + + const line = lineOrSelections + const activeOnly = options?.activeOnly ?? true + + for (const selection of this._selections) { + if ( + line === selection.active || + (!activeOnly && + ((selection.anchor >= line && line >= selection.active) || + (selection.active >= line && line >= selection.anchor))) + ) { + return true + } + } + return false + } +} + +function isIncluded(selections: LineSelection[] | undefined, within: LineSelection[] | undefined): boolean { + if (selections === undefined && within === undefined) { + return true + } + if (selections === undefined || within === undefined || selections.length !== within.length) { + return false + } + + return selections.every((s, i) => { + const match = within[i] + return s.active === match.active && s.anchor === match.anchor + }) +} + +function toLineSelections(selections: readonly vscode.Selection[]): LineSelection[] +function toLineSelections(selections: readonly vscode.Selection[] | undefined): LineSelection[] | undefined +function toLineSelections(selections: readonly vscode.Selection[] | undefined) { + return selections?.map((s) => ({ active: s.active.line, anchor: s.anchor.line })) +} diff --git a/packages/amazonq/src/app/inline/telemetryHelper.ts b/packages/amazonq/src/app/inline/telemetryHelper.ts new file mode 100644 index 00000000000..2758ffe4c13 --- /dev/null +++ b/packages/amazonq/src/app/inline/telemetryHelper.ts @@ -0,0 +1,162 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import { AuthUtil, getSelectedCustomization } from 'aws-core-vscode/codewhisperer' +import { CodewhispererLanguage } from 'aws-core-vscode/shared' +import { CodewhispererTriggerType, telemetry } from 'aws-core-vscode/telemetry' +import { InlineCompletionTriggerKind } from 'vscode' + +export class TelemetryHelper { + // Variables needed for client component latency + private _invokeSuggestionStartTime = 0 + private _preprocessEndTime = 0 + private _sdkApiCallStartTime = 0 + private _sdkApiCallEndTime = 0 + private _allPaginationEndTime = 0 + private _firstSuggestionShowTime = 0 + private _firstResponseRequestId = '' + private _sessionId = '' + private _language: CodewhispererLanguage = 'java' + private _triggerType: CodewhispererTriggerType = 'OnDemand' + + constructor() {} + + static #instance: TelemetryHelper + + public static get instance() { + return (this.#instance ??= new this()) + } + + public resetClientComponentLatencyTime() { + this._invokeSuggestionStartTime = 0 + this._preprocessEndTime = 0 + this._sdkApiCallStartTime = 0 + this._sdkApiCallEndTime = 0 + this._firstSuggestionShowTime = 0 + this._allPaginationEndTime = 0 + this._firstResponseRequestId = '' + } + + public setInvokeSuggestionStartTime() { + this.resetClientComponentLatencyTime() + this._invokeSuggestionStartTime = performance.now() + } + + get invokeSuggestionStartTime(): number { + return this._invokeSuggestionStartTime + } + + public setPreprocessEndTime() { + this._preprocessEndTime = performance.now() + } + + get preprocessEndTime(): number { + return this._preprocessEndTime + } + + public setSdkApiCallStartTime() { + if (this._sdkApiCallStartTime === 0) { + this._sdkApiCallStartTime = performance.now() + } + } + + get sdkApiCallStartTime(): number { + return this._sdkApiCallStartTime + } + + public setSdkApiCallEndTime() { + if (this._sdkApiCallEndTime === 0 && this._sdkApiCallStartTime !== 0) { + this._sdkApiCallEndTime = performance.now() + } + } + + get sdkApiCallEndTime(): number { + return this._sdkApiCallEndTime + } + + public setAllPaginationEndTime() { + if (this._allPaginationEndTime === 0 && this._sdkApiCallEndTime !== 0) { + this._allPaginationEndTime = performance.now() + } + } + + get allPaginationEndTime(): number { + return this._allPaginationEndTime + } + + public setFirstSuggestionShowTime() { + if (this._firstSuggestionShowTime === 0 && this._sdkApiCallEndTime !== 0) { + this._firstSuggestionShowTime = performance.now() + } + } + + get firstSuggestionShowTime(): number { + return this._firstSuggestionShowTime + } + + public setFirstResponseRequestId(requestId: string) { + if (this._firstResponseRequestId === '') { + this._firstResponseRequestId = requestId + } + } + + get firstResponseRequestId(): string { + return this._firstResponseRequestId + } + + public setSessionId(sessionId: string) { + if (this._sessionId === '') { + this._sessionId = sessionId + } + } + + get sessionId(): string { + return this._sessionId + } + + public setLanguage(language: CodewhispererLanguage) { + this._language = language + } + + get language(): CodewhispererLanguage { + return this._language + } + + public setTriggerType(triggerType: InlineCompletionTriggerKind) { + if (triggerType === InlineCompletionTriggerKind.Invoke) { + this._triggerType = 'OnDemand' + } else if (triggerType === InlineCompletionTriggerKind.Automatic) { + this._triggerType = 'AutoTrigger' + } + } + + get triggerType(): string { + return this._triggerType + } + + // report client component latency after all pagination call finish + // and at least one suggestion is shown to the user + public tryRecordClientComponentLatency() { + if (this._firstSuggestionShowTime === 0 || this._allPaginationEndTime === 0) { + return + } + telemetry.codewhisperer_clientComponentLatency.emit({ + codewhispererAllCompletionsLatency: this._allPaginationEndTime - this._sdkApiCallStartTime, + codewhispererCompletionType: 'Line', + codewhispererCredentialFetchingLatency: 0, // no longer relevant, because we don't re-build the sdk. Flare already has that set + codewhispererCustomizationArn: getSelectedCustomization().arn, + codewhispererEndToEndLatency: this._firstSuggestionShowTime - this._invokeSuggestionStartTime, + codewhispererFirstCompletionLatency: this._sdkApiCallEndTime - this._sdkApiCallStartTime, + codewhispererLanguage: this._language, + codewhispererPostprocessingLatency: this._firstSuggestionShowTime - this._sdkApiCallEndTime, + codewhispererPreprocessingLatency: this._preprocessEndTime - this._invokeSuggestionStartTime, + codewhispererRequestId: this._firstResponseRequestId, + codewhispererSessionId: this._sessionId, + codewhispererTriggerType: this._triggerType, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, + result: 'Succeeded', + }) + } +} diff --git a/packages/amazonq/src/inlineChat/decorations/inlineLineAnnotationController.ts b/packages/amazonq/src/app/inline/tutorials/inlineChatTutorialAnnotation.ts similarity index 72% rename from packages/amazonq/src/inlineChat/decorations/inlineLineAnnotationController.ts rename to packages/amazonq/src/app/inline/tutorials/inlineChatTutorialAnnotation.ts index 9ec5e08122d..1208b4766af 100644 --- a/packages/amazonq/src/inlineChat/decorations/inlineLineAnnotationController.ts +++ b/packages/amazonq/src/app/inline/tutorials/inlineChatTutorialAnnotation.ts @@ -3,14 +3,15 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { Container } from 'aws-core-vscode/codewhisperer' import * as vscode from 'vscode' +import { InlineTutorialAnnotation } from './inlineTutorialAnnotation' +import { globals } from 'aws-core-vscode/shared' -export class InlineLineAnnotationController { +export class InlineChatTutorialAnnotation { private enabled: boolean = true - constructor(context: vscode.ExtensionContext) { - context.subscriptions.push( + constructor(private readonly inlineTutorialAnnotation: InlineTutorialAnnotation) { + globals.context.subscriptions.push( vscode.window.onDidChangeTextEditorSelection(async ({ selections, textEditor }) => { let showShow = false @@ -33,12 +34,12 @@ export class InlineLineAnnotationController { private async setVisible(editor: vscode.TextEditor, visible: boolean) { let needsRefresh: boolean if (visible) { - needsRefresh = await Container.instance.lineAnnotationController.tryShowInlineHint() + needsRefresh = await this.inlineTutorialAnnotation.tryShowInlineHint() } else { - needsRefresh = await Container.instance.lineAnnotationController.tryHideInlineHint() + needsRefresh = await this.inlineTutorialAnnotation.tryHideInlineHint() } if (needsRefresh) { - await Container.instance.lineAnnotationController.refresh(editor, 'codewhisperer') + await this.inlineTutorialAnnotation.refresh(editor, 'codewhisperer') } } diff --git a/packages/amazonq/src/app/inline/tutorials/inlineTutorialAnnotation.ts b/packages/amazonq/src/app/inline/tutorials/inlineTutorialAnnotation.ts new file mode 100644 index 00000000000..7fd8ace5d94 --- /dev/null +++ b/packages/amazonq/src/app/inline/tutorials/inlineTutorialAnnotation.ts @@ -0,0 +1,521 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as vscode from 'vscode' +import * as os from 'os' +import { + AnnotationChangeSource, + AuthUtil, + inlinehintKey, + runtimeLanguageContext, + TelemetryHelper, +} from 'aws-core-vscode/codewhisperer' +import { editorUtilities, getLogger, globals, setContext, vscodeUtilities } from 'aws-core-vscode/shared' +import { LinesChangeEvent, LineSelection, LineTracker } from '../stateTracker/lineTracker' +import { telemetry } from 'aws-core-vscode/telemetry' +import { cancellableDebounce } from 'aws-core-vscode/utils' +import { SessionManager } from '../sessionManager' + +const case3TimeWindow = 30000 // 30 seconds + +const maxSmallIntegerV8 = 2 ** 30 // Max number that can be stored in V8's smis (small integers) + +function fromId(id: string | undefined, sessionManager: SessionManager): AnnotationState | undefined { + switch (id) { + case AutotriggerState.id: + return new AutotriggerState(sessionManager) + case PressTabState.id: + return new AutotriggerState(sessionManager) + case ManualtriggerState.id: + return new ManualtriggerState() + case TryMoreExState.id: + return new TryMoreExState() + case EndState.id: + return new EndState() + case InlineChatState.id: + return new InlineChatState() + default: + return undefined + } +} + +interface AnnotationState { + id: string + suppressWhileRunning: boolean + decorationRenderOptions?: vscode.ThemableDecorationAttachmentRenderOptions + + text: () => string + updateState(changeSource: AnnotationChangeSource, force: boolean): AnnotationState | undefined + isNextState(state: AnnotationState | undefined): boolean +} + +/** + * case 1: How Cwspr triggers + * Trigger Criteria: + * User opens an editor file && + * CW is not providing a suggestion && + * User has not accepted any suggestion + * + * Exit criteria: + * User accepts 1 suggestion + * + */ +export class AutotriggerState implements AnnotationState { + static id = 'codewhisperer_learnmore_case_1' + id = AutotriggerState.id + + suppressWhileRunning = true + text = () => 'Amazon Q Tip 1/3: Start typing to get suggestions ([ESC] to exit)' + static acceptedCount = 0 + + constructor(private readonly sessionManager: SessionManager) {} + + updateState(changeSource: AnnotationChangeSource, force: boolean): AnnotationState | undefined { + if (AutotriggerState.acceptedCount < this.sessionManager.acceptedSuggestionCount) { + return new ManualtriggerState() + } else if (this.sessionManager.getActiveRecommendation().length > 0) { + return new PressTabState(this.sessionManager) + } else { + return this + } + } + + isNextState(state: AnnotationState | undefined): boolean { + return state instanceof ManualtriggerState + } +} + +/** + * case 1-a: Tab to accept + * Trigger Criteria: + * Case 1 && + * Inline suggestion is being shown + * + * Exit criteria: + * User accepts 1 suggestion + */ +export class PressTabState implements AnnotationState { + static id = 'codewhisperer_learnmore_case_1a' + id = PressTabState.id + + suppressWhileRunning = false + + text = () => 'Amazon Q Tip 1/3: Press [TAB] to accept ([ESC] to exit)' + + constructor(private readonly sessionManager: SessionManager) {} + + updateState(changeSource: AnnotationChangeSource, force: boolean): AnnotationState | undefined { + return new AutotriggerState(this.sessionManager).updateState(changeSource, force) + } + + isNextState(state: AnnotationState | undefined): boolean { + return state instanceof ManualtriggerState + } +} + +/** + * case 2: Manual trigger + * Trigger Criteria: + * User exists case 1 && + * User navigates to a new line + * + * Exit criteria: + * User inokes manual trigger shortcut + */ +export class ManualtriggerState implements AnnotationState { + static id = 'codewhisperer_learnmore_case_2' + id = ManualtriggerState.id + + suppressWhileRunning = true + + text = () => { + if (os.platform() === 'win32') { + return 'Amazon Q Tip 2/3: Invoke suggestions with [Alt] + [C] ([ESC] to exit)' + } + + return 'Amazon Q Tip 2/3: Invoke suggestions with [Option] + [C] ([ESC] to exit)' + } + hasManualTrigger: boolean = false + hasValidResponse: boolean = false + + updateState(changeSource: AnnotationChangeSource, force: boolean): AnnotationState | undefined { + if (this.hasManualTrigger && this.hasValidResponse) { + if (changeSource !== 'codewhisperer') { + return new TryMoreExState() + } else { + return undefined + } + } else { + return this + } + } + + isNextState(state: AnnotationState | undefined): boolean { + return state instanceof TryMoreExState + } +} + +/** + * case 3: Learn more + * Trigger Criteria: + * User exists case 2 && + * User navigates to a new line + * + * Exit criteria: + * User accepts or rejects the suggestion + */ +export class TryMoreExState implements AnnotationState { + static id = 'codewhisperer_learnmore_case_3' + id = TryMoreExState.id + + suppressWhileRunning = true + + text = () => 'Amazon Q Tip 3/3: For settings, open the Amazon Q menu from the status bar ([ESC] to exit)' + updateState(changeSource: AnnotationChangeSource, force: boolean): AnnotationState { + if (force) { + return new EndState() + } + return this + } + + isNextState(state: AnnotationState | undefined): boolean { + return state instanceof EndState + } + + static learnmoeCount: number = 0 +} + +export class EndState implements AnnotationState { + static id = 'codewhisperer_learnmore_end' + id = EndState.id + + suppressWhileRunning = true + text = () => '' + updateState(changeSource: AnnotationChangeSource, force: boolean): AnnotationState { + return this + } + isNextState(state: AnnotationState): boolean { + return false + } +} + +export class InlineChatState implements AnnotationState { + static id = 'amazonq_annotation_inline_chat' + id = InlineChatState.id + suppressWhileRunning = false + + text = () => { + if (os.platform() === 'darwin') { + return 'Amazon Q: Edit \u2318I' + } + return 'Amazon Q: Edit (Ctrl+I)' + } + updateState(_changeSource: AnnotationChangeSource, _force: boolean): AnnotationState { + return this + } + isNextState(_state: AnnotationState | undefined): boolean { + return false + } +} + +/** + * There are + * - existing users + * - new users + * -- new users who has not seen tutorial + * -- new users who has seen tutorial + * + * "existing users" should have the context key "CODEWHISPERER_AUTO_TRIGGER_ENABLED" + * "new users who has seen tutorial" should have the context key "inlineKey" and "CODEWHISPERER_AUTO_TRIGGER_ENABLED" + * the remaining grouop of users should belong to "new users who has not seen tutorial" + */ +export class InlineTutorialAnnotation implements vscode.Disposable { + private readonly _disposable: vscode.Disposable + private _editor: vscode.TextEditor | undefined + + private _currentState: AnnotationState + + private readonly cwLineHintDecoration: vscode.TextEditorDecorationType = + vscode.window.createTextEditorDecorationType({ + after: { + margin: '0 0 0 3em', + // "borderRadius" and "padding" are not available on "after" type of decoration, this is a hack to inject these css prop to "after" content. Refer to https://github.com/microsoft/vscode/issues/68845 + textDecoration: ';border-radius:0.25rem;padding:0rem 0.5rem;', + width: 'fit-content', + }, + rangeBehavior: vscode.DecorationRangeBehavior.OpenOpen, + }) + + constructor( + private readonly lineTracker: LineTracker, + private readonly sessionManager: SessionManager + ) { + const cachedState = fromId(globals.globalState.get(inlinehintKey), sessionManager) + const cachedAutotriggerEnabled = globals.globalState.get('CODEWHISPERER_AUTO_TRIGGER_ENABLED') + + // new users (has or has not seen tutorial) + if (cachedAutotriggerEnabled === undefined || cachedState !== undefined) { + this._currentState = cachedState ?? new AutotriggerState(this.sessionManager) + getLogger().debug( + `codewhisperer: new user login, activating inline tutorial. (autotriggerEnabled=${cachedAutotriggerEnabled}; inlineState=${cachedState?.id})` + ) + } else { + this._currentState = new EndState() + getLogger().debug(`codewhisperer: existing user login, disabling inline tutorial.`) + } + + this._disposable = vscode.Disposable.from( + vscodeUtilities.subscribeOnce(this.lineTracker.onReady)(async (_) => { + await this.onReady() + }), + this.lineTracker.onDidChangeActiveLines(async (e) => { + await this.onActiveLinesChanged(e) + }), + AuthUtil.instance.onDidChangeConnectionState(async (e) => { + await this.refresh(vscode.window.activeTextEditor, 'editor') + }) + ) + } + + dispose() { + this._disposable.dispose() + } + + private _isReady: boolean = false + + private async onReady(): Promise { + this._isReady = !(this._currentState instanceof EndState) + await this._refresh(vscode.window.activeTextEditor, 'editor') + } + + async triggered(triggerType: vscode.InlineCompletionTriggerKind): Promise { + await telemetry.withTraceId(async () => { + if (!this._isReady) { + return + } + + if (this._currentState instanceof ManualtriggerState) { + if ( + triggerType === vscode.InlineCompletionTriggerKind.Invoke && + this._currentState.hasManualTrigger === false + ) { + this._currentState.hasManualTrigger = true + } + if ( + this.sessionManager.getActiveRecommendation().length > 0 && + this._currentState.hasValidResponse === false + ) { + this._currentState.hasValidResponse = true + } + } + + await this.refresh(vscode.window.activeTextEditor, 'codewhisperer') + }, TelemetryHelper.instance.traceId) + } + + isTutorialDone(): boolean { + return this._currentState.id === new EndState().id + } + + isInlineChatHint(): boolean { + return this._currentState.id === new InlineChatState().id + } + + async dismissTutorial() { + this._currentState = new EndState() + await setContext('aws.codewhisperer.tutorial.workInProgress', false) + await globals.globalState.update(inlinehintKey, this._currentState.id) + } + + /** + * Trys to show the inline hint, if the tutorial is not finished it will not be shown + */ + async tryShowInlineHint(): Promise { + if (this.isTutorialDone()) { + this._isReady = true + this._currentState = new InlineChatState() + return true + } + return false + } + + async tryHideInlineHint(): Promise { + if (this._currentState instanceof InlineChatState) { + this._currentState = new EndState() + return true + } + return false + } + + private async onActiveLinesChanged(e: LinesChangeEvent) { + if (!this._isReady) { + return + } + + this.clear() + + await this.refresh(e.editor, e.reason) + } + + clear() { + this._editor?.setDecorations(this.cwLineHintDecoration, []) + } + + async refresh(editor: vscode.TextEditor | undefined, source: AnnotationChangeSource, force?: boolean) { + if (force) { + this.refreshDebounced.cancel() + await this._refresh(editor, source, true) + } else { + await this.refreshDebounced.promise(editor, source) + } + } + + private readonly refreshDebounced = cancellableDebounce( + async (editor: vscode.TextEditor | undefined, source: AnnotationChangeSource, force?: boolean) => { + await this._refresh(editor, source, force) + }, + 250 + ) + + private async _refresh(editor: vscode.TextEditor | undefined, source: AnnotationChangeSource, force?: boolean) { + if (!this._isReady) { + this.clear() + return + } + + if (this.isTutorialDone()) { + this.clear() + return + } + + if (editor === undefined && this._editor === undefined) { + this.clear() + return + } + + const selections = this.lineTracker.selections + if (editor === undefined || selections === undefined || !editorUtilities.isTextEditor(editor)) { + this.clear() + return + } + + if (this._editor !== editor) { + // Clear any annotations on the previously active editor + this.clear() + this._editor = editor + } + + // Make sure the editor hasn't died since the await above and that we are still on the same line(s) + if (editor.document === undefined || !this.lineTracker.includes(selections)) { + this.clear() + return + } + + if (!AuthUtil.instance.isConnected()) { + this.clear() + return + } + + // Disable Tips when language is not supported by Amazon Q. + if (!runtimeLanguageContext.isLanguageSupported(editor.document)) { + return + } + + await this.updateDecorations(editor, selections, source, force) + } + + private async updateDecorations( + editor: vscode.TextEditor, + lines: LineSelection[], + source: AnnotationChangeSource, + force?: boolean + ) { + const range = editor.document.validateRange( + new vscode.Range(lines[0].active, maxSmallIntegerV8, lines[0].active, maxSmallIntegerV8) + ) + + const decorationOptions = this.getInlineDecoration(editor, lines, source, force) as + | vscode.DecorationOptions + | undefined + + if (decorationOptions === undefined) { + this.clear() + await setContext('aws.codewhisperer.tutorial.workInProgress', false) + return + } else if (this.isTutorialDone()) { + // special case + // Endstate is meaningless and doesnt need to be rendered + this.clear() + await this.dismissTutorial() + return + } else if (decorationOptions.renderOptions?.after?.contentText === new TryMoreExState().text()) { + // special case + // case 3 exit criteria is to fade away in 30s + setTimeout(async () => { + await this.refresh(editor, source, true) + }, case3TimeWindow) + } + + decorationOptions.range = range + + await globals.globalState.update(inlinehintKey, this._currentState.id) + if (!this.isInlineChatHint()) { + await setContext('aws.codewhisperer.tutorial.workInProgress', true) + } + editor.setDecorations(this.cwLineHintDecoration, [decorationOptions]) + } + + getInlineDecoration( + editor: vscode.TextEditor, + lines: LineSelection[], + source: AnnotationChangeSource, + force?: boolean + ): Partial | undefined { + const isCWRunning = this.sessionManager.getActiveSession()?.isRequestInProgress ?? false + + const textOptions: vscode.ThemableDecorationAttachmentRenderOptions = { + contentText: '', + fontWeight: 'normal', + fontStyle: 'normal', + textDecoration: 'none', + color: 'var(--vscode-editor-background)', + backgroundColor: 'var(--vscode-foreground)', + } + + if (isCWRunning && this._currentState.suppressWhileRunning) { + return undefined + } + + const updatedState: AnnotationState | undefined = this._currentState.updateState(source, force ?? false) + + if (updatedState === undefined) { + return undefined + } + + if (this._currentState.isNextState(updatedState)) { + // special case because PressTabState is part of case_1 (1a) which possibly jumps directly from case_1a to case_2 and miss case_1 + if (this._currentState instanceof PressTabState) { + telemetry.ui_click.emit({ elementId: AutotriggerState.id, passive: true }) + } + telemetry.ui_click.emit({ elementId: this._currentState.id, passive: true }) + } + + // update state + this._currentState = updatedState + + // take snapshot of accepted session so that we can compre if there is delta -> users accept 1 suggestion after seeing this state + AutotriggerState.acceptedCount = this.sessionManager.acceptedSuggestionCount + + textOptions.contentText = this._currentState.text() + + return { + renderOptions: { after: textOptions }, + } + } + + public get currentState(): AnnotationState { + return this._currentState + } +} diff --git a/packages/amazonq/src/extension.ts b/packages/amazonq/src/extension.ts index 45641b37440..b034f360cec 100644 --- a/packages/amazonq/src/extension.ts +++ b/packages/amazonq/src/extension.ts @@ -3,8 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { Auth, AuthUtils, CredentialsStore, LoginManager, initializeAuth } from 'aws-core-vscode/auth' -import { activate as activateCodeWhisperer, shutdown as shutdownCodeWhisperer } from 'aws-core-vscode/codewhisperer' +import { authUtils, CredentialsStore, LoginManager } from 'aws-core-vscode/auth' +import { + activate as activateCodeWhisperer, + refreshStatusBar, + shutdown as shutdownCodeWhisperer, + updateReferenceLog, +} from 'aws-core-vscode/codewhisperer' import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode' import { CommonAuthWebview } from 'aws-core-vscode/login' import { @@ -33,7 +38,7 @@ import { maybeShowMinVscodeWarning, Experiments, isSageMaker, - isAmazonLinux2, + Commands, } from 'aws-core-vscode/shared' import { ExtStartUpSources } from 'aws-core-vscode/telemetry' import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils' @@ -44,7 +49,6 @@ import { registerCommands } from './commands' import { focusAmazonQPanel } from 'aws-core-vscode/codewhispererChat' import { activate as activateAmazonqLsp } from './lsp/activation' import { activate as activateInlineCompletion } from './app/inline/activation' -import { hasGlibcPatch } from './lsp/client' export const amazonQContextPrefix = 'amazonq' @@ -114,21 +118,17 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is await activateTelemetry(context, globals.awsContext, Settings.instance, 'Amazon Q For VS Code') - await initializeAuth(globals.loginManager) - const extContext = { extensionContext: context, } - // This contains every lsp agnostic things (auth, security scan, code scan) + + activateAuthDependentCommands() + + // Auth is dependent on LSP, needs to be activated before CW and Inline + await activateAmazonqLsp(context) + + // This contains every lsp agnostic things (security scan, code scan) await activateCodeWhisperer(extContext as ExtContext) - if ( - (Experiments.instance.get('amazonqLSP', true) || Auth.instance.isInternalAmazonUser()) && - (!isAmazonLinux2() || hasGlibcPatch()) - ) { - // start the Amazon Q LSP for internal users first - // for AL2, start LSP if glibc patch is found - await activateAmazonqLsp(context) - } if (!Experiments.instance.get('amazonqLSPInline', false)) { await activateInlineCompletion() } @@ -136,6 +136,10 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is // Generic extension commands registerGenericCommands(context, amazonQContextPrefix) + // Create status bar and reference log UI elements + void Commands.tryExecute('aws.amazonq.refreshStatusBar') + void Commands.tryExecute('aws.amazonq.updateReferenceLog') + // Amazon Q specific commands registerCommands(context) @@ -162,7 +166,7 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is // reload webviews await vscode.commands.executeCommand('workbench.action.webview.reloadWebviewAction') - if (AuthUtils.ExtensionUse.instance.isFirstUse()) { + if (authUtils.ExtensionUse.instance.isFirstUse()) { // Give time for the extension to finish initializing. globals.clock.setTimeout(async () => { CommonAuthWebview.authSource = ExtStartUpSources.firstStartUp @@ -172,7 +176,7 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is context.subscriptions.push( Experiments.instance.onDidChange(async (event) => { - if (event.key === 'amazonqLSP' || event.key === 'amazonqChatLSP' || event.key === 'amazonqLSPInline') { + if (event.key === 'amazonqChatLSP' || event.key === 'amazonqLSPInline') { await vscode.window .showInformationMessage( 'Amazon Q LSP setting has changed. Reload VS Code for the changes to take effect.', @@ -186,6 +190,14 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is } }) ) + + // Activate commands that are required for activateAmazonqLsp + function activateAuthDependentCommands() { + // update reference log instance + updateReferenceLog.register() + // refresh codewhisperer status bar + refreshStatusBar.register() + } } export async function deactivateCommon() { diff --git a/packages/amazonq/src/extensionNode.ts b/packages/amazonq/src/extensionNode.ts index 8224b9ce310..576757c36e2 100644 --- a/packages/amazonq/src/extensionNode.ts +++ b/packages/amazonq/src/extensionNode.ts @@ -7,28 +7,19 @@ import * as vscode from 'vscode' import { activateAmazonQCommon, amazonQContextPrefix, deactivateCommon } from './extension' import { DefaultAmazonQAppInitContext, AmazonQChatViewProvider } from 'aws-core-vscode/amazonq' import { activate as activateTransformationHub } from 'aws-core-vscode/amazonqGumby' -import { - ExtContext, - globals, - CrashMonitoring, - getLogger, - isNetworkError, - isSageMaker, - Experiments, -} from 'aws-core-vscode/shared' +import { ExtContext, globals, CrashMonitoring, getLogger, isSageMaker, Experiments } from 'aws-core-vscode/shared' import { filetypes, SchemaService } from 'aws-core-vscode/sharedNode' import { updateDevMode } from 'aws-core-vscode/dev' import { CommonAuthViewProvider } from 'aws-core-vscode/login' import { isExtensionActive, VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils' import { registerSubmitFeedback } from 'aws-core-vscode/feedback' import { DevOptions } from 'aws-core-vscode/dev' -import { Auth, AuthUtils, getTelemetryMetadataForConn, isAnySsoConnection } from 'aws-core-vscode/auth' +import { Auth, authUtils } from 'aws-core-vscode/auth' import api from './api' import { activate as activateCWChat } from './app/chat/activation' -import { activate as activateInlineChat } from './inlineChat/activation' import { beta } from 'aws-core-vscode/dev' import { activate as activateNotifications, NotificationsController } from 'aws-core-vscode/notifications' -import { AuthState, AuthUtil } from 'aws-core-vscode/codewhisperer' +import { AuthUtil } from 'aws-core-vscode/codewhisperer' import { telemetry, AuthUserState } from 'aws-core-vscode/telemetry' import { activateAgents } from './app/chat/node/activateAgents' @@ -73,7 +64,6 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) { } activateAgents() await activateTransformationHub(extContext as ExtContext) - activateInlineChat(context) const authProvider = new CommonAuthViewProvider( context, @@ -95,52 +85,30 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) { await setupDevMode(context) await beta.activate(context) - // TODO: Should probably emit for web as well. - // Will the web metric look the same? + await getAuthState() telemetry.auth_userState.emit({ passive: true, result: 'Succeeded', - source: AuthUtils.ExtensionUse.instance.sourceForTelemetry(), - ...(await getAuthState()), + source: authUtils.ExtensionUse.instance.sourceForTelemetry(), + authStatus: AuthUtil.instance.getAuthState(), + authEnabledConnections: (await AuthUtil.instance.getAuthFormIds()).join(','), }) void activateNotifications(context, getAuthState) } async function getAuthState(): Promise> { - 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) - } + const state = AuthUtil.instance.getAuthState() - // Pending profile selection state means users already log in with Sso service - if (authState === 'pendingProfileSelection') { - authState = 'connected' + if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || isSageMaker())) { + getLogger().error('Current Amazon Q connection is not SSO') } return { - authStatus: - authState === 'connected' || authState === 'expired' || authState === 'connectedWithNetworkError' - ? authState - : 'notConnected', - authEnabledConnections: AuthUtils.getAuthFormIdsFromConnection(currConn).join(','), - ...(await getTelemetryMetadataForConn(currConn)), + // @ts-ignore + authStatus: (state ?? 'notConnected') as AuthStatus, + authEnabledConnections: (await AuthUtil.instance.getAuthFormIds()).join(','), + ...AuthUtil.instance.getTelemetryMetadata(), } } diff --git a/packages/amazonq/src/inlineChat/activation.ts b/packages/amazonq/src/inlineChat/activation.ts index a42dfdb3e02..9f196f31ba3 100644 --- a/packages/amazonq/src/inlineChat/activation.ts +++ b/packages/amazonq/src/inlineChat/activation.ts @@ -5,8 +5,15 @@ import * as vscode from 'vscode' import { InlineChatController } from './controller/inlineChatController' import { registerInlineCommands } from './command/registerInlineCommands' +import { LanguageClient } from 'vscode-languageclient' +import { InlineChatTutorialAnnotation } from '../app/inline/tutorials/inlineChatTutorialAnnotation' -export function activate(context: vscode.ExtensionContext) { - const inlineChatController = new InlineChatController(context) +export function activate( + context: vscode.ExtensionContext, + client: LanguageClient, + encryptionKey: Buffer, + inlineChatTutorialAnnotation: InlineChatTutorialAnnotation +) { + const inlineChatController = new InlineChatController(context, client, encryptionKey, inlineChatTutorialAnnotation) registerInlineCommands(context, inlineChatController) } diff --git a/packages/amazonq/src/inlineChat/controller/inlineChatController.ts b/packages/amazonq/src/inlineChat/controller/inlineChatController.ts index 7ace8d0095e..7151a8f9723 100644 --- a/packages/amazonq/src/inlineChat/controller/inlineChatController.ts +++ b/packages/amazonq/src/inlineChat/controller/inlineChatController.ts @@ -14,6 +14,7 @@ import { CodelensProvider } from '../codeLenses/codeLenseProvider' import { PromptMessage, ReferenceLogController } from 'aws-core-vscode/codewhispererChat' import { CodeWhispererSettings } from 'aws-core-vscode/codewhisperer' import { UserWrittenCodeTracker } from 'aws-core-vscode/codewhisperer' +import { LanguageClient } from 'vscode-languageclient' import { codicon, getIcon, @@ -23,8 +24,9 @@ import { Timeout, textDocumentUtil, isSageMaker, + Experiments, } from 'aws-core-vscode/shared' -import { InlineLineAnnotationController } from '../decorations/inlineLineAnnotationController' +import { InlineChatTutorialAnnotation } from '../../app/inline/tutorials/inlineChatTutorialAnnotation' export class InlineChatController { private task: InlineTask | undefined @@ -32,15 +34,24 @@ export class InlineChatController { private readonly inlineChatProvider: InlineChatProvider private readonly codeLenseProvider: CodelensProvider private readonly referenceLogController = new ReferenceLogController() - private readonly inlineLineAnnotationController: InlineLineAnnotationController + private readonly inlineChatTutorialAnnotation: InlineChatTutorialAnnotation + private readonly computeDiffAndRenderOnEditor: (query: string) => Promise private userQuery: string | undefined private listeners: vscode.Disposable[] = [] - constructor(context: vscode.ExtensionContext) { - this.inlineChatProvider = new InlineChatProvider() + constructor( + context: vscode.ExtensionContext, + client: LanguageClient, + encryptionKey: Buffer, + inlineChatTutorialAnnotation: InlineChatTutorialAnnotation + ) { + this.inlineChatProvider = new InlineChatProvider(client, encryptionKey) this.inlineChatProvider.onErrorOccured(() => this.handleError()) this.codeLenseProvider = new CodelensProvider(context) - this.inlineLineAnnotationController = new InlineLineAnnotationController(context) + this.inlineChatTutorialAnnotation = inlineChatTutorialAnnotation + this.computeDiffAndRenderOnEditor = Experiments.instance.get('amazonqLSPInlineChat', false) + ? this.computeDiffAndRenderOnEditorLSP.bind(this) + : this.computeDiffAndRenderOnEditorLocal.bind(this) } public async createTask( @@ -138,7 +149,7 @@ export class InlineChatController { this.codeLenseProvider.updateLenses(task) if (task.state === TaskState.InProgress) { if (vscode.window.activeTextEditor) { - await this.inlineLineAnnotationController.hide(vscode.window.activeTextEditor) + await this.inlineChatTutorialAnnotation.hide(vscode.window.activeTextEditor) } } await this.refreshCodeLenses(task) @@ -164,7 +175,7 @@ export class InlineChatController { this.listeners = [] this.task = undefined - this.inlineLineAnnotationController.enable() + this.inlineChatTutorialAnnotation.enable() await setContext('amazonq.inline.codelensShortcutEnabled', undefined) } @@ -205,8 +216,8 @@ export class InlineChatController { this.userQuery = query await textDocumentUtil.addEofNewline(editor) this.task = await this.createTask(query, editor.document, editor.selection) - await this.inlineLineAnnotationController.disable(editor) - await this.computeDiffAndRenderOnEditor(query, editor.document).catch(async (err) => { + await this.inlineChatTutorialAnnotation.disable(editor) + await this.computeDiffAndRenderOnEditor(query).catch(async (err) => { getLogger().error('computeDiffAndRenderOnEditor error: %s', (err as Error)?.message) if (err instanceof Error) { void vscode.window.showErrorMessage(`Amazon Q: ${err.message}`) @@ -218,7 +229,46 @@ export class InlineChatController { }) } - private async computeDiffAndRenderOnEditor(query: string, document: vscode.TextDocument) { + private async computeDiffAndRenderOnEditorLSP(query: string) { + if (!this.task) { + return + } + + await this.updateTaskAndLenses(this.task, TaskState.InProgress) + getLogger().info(`inline chat query:\n${query}`) + const uuid = randomUUID() + const message: PromptMessage = { + message: query, + messageId: uuid, + command: undefined, + userIntent: undefined, + tabID: uuid, + } + + const response = await this.inlineChatProvider.processPromptMessageLSP(message) + + // TODO: add tests for this case. + if (!response.body) { + getLogger().warn('Empty body in inline chat response') + await this.handleError() + return + } + + // Update inline diff view + const textDiff = computeDiff(response.body, this.task, false) + const decorations = computeDecorations(this.task) + this.task.decorations = decorations + await this.applyDiff(this.task, textDiff ?? []) + this.decorator.applyDecorations(this.task) + + // Update Codelenses + await this.updateTaskAndLenses(this.task, TaskState.WaitingForDecision) + await setContext('amazonq.inline.codelensShortcutEnabled', true) + this.undoListener(this.task) + } + + // TODO: remove this implementation in favor of LSP + private async computeDiffAndRenderOnEditorLocal(query: string) { if (!this.task) { return } diff --git a/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts b/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts index e6534d65532..64a67224a2e 100644 --- a/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts +++ b/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts @@ -8,6 +8,8 @@ import { CodeWhispererStreamingServiceException, GenerateAssistantResponseCommandOutput, } from '@amzn/codewhisperer-streaming' +import { LanguageClient } from 'vscode-languageclient' +import { inlineChatRequestType } from '@aws/language-server-runtimes/protocol' import { AuthUtil, getSelectedCustomization } from 'aws-core-vscode/codewhisperer' import { ChatSessionStorage, @@ -25,6 +27,9 @@ import { codeWhispererClient } from 'aws-core-vscode/codewhisperer' import type { InlineChatEvent } from 'aws-core-vscode/codewhisperer' import { InlineTask } from '../controller/inlineTask' import { extractAuthFollowUp } from 'aws-core-vscode/amazonq' +import { InlineChatParams, InlineChatResult } from '@aws/language-server-runtimes-types' +import { decryptResponse, encryptRequest } from '../../lsp/encryption' +import { getCursorState } from '../../lsp/utils' export class InlineChatProvider { private readonly editorContextExtractor: EditorContextExtractor @@ -34,13 +39,49 @@ export class InlineChatProvider { private errorEmitter = new vscode.EventEmitter() public onErrorOccured = this.errorEmitter.event - public constructor() { + public constructor( + private readonly client: LanguageClient, + private readonly encryptionKey: Buffer + ) { this.editorContextExtractor = new EditorContextExtractor() this.userIntentRecognizer = new UserIntentRecognizer() this.sessionStorage = new ChatSessionStorage() this.triggerEventsStorage = new TriggerEventsStorage() } + private getCurrentEditorParams(prompt: string): InlineChatParams { + const editor = vscode.window.activeTextEditor + if (!editor) { + throw new ToolkitError('No active editor') + } + + const documentUri = editor.document.uri.toString() + const cursorState = getCursorState(editor.selections) + return { + prompt: { + prompt, + }, + cursorState, + textDocument: { + uri: documentUri, + }, + } + } + + public async processPromptMessageLSP(message: PromptMessage): Promise { + // TODO: handle partial responses. + getLogger().info('Making inline chat request with message %O', message) + const params = this.getCurrentEditorParams(message.message ?? '') + + const inlineChatRequest = await encryptRequest(params, this.encryptionKey) + const response = await this.client.sendRequest(inlineChatRequestType.method, inlineChatRequest) + const inlineChatResponse = await decryptResponse(response, this.encryptionKey) + this.client.info(`Logging response for inline chat ${JSON.stringify(inlineChatResponse)}`) + + return inlineChatResponse + } + + // TODO: remove in favor of LSP implementation. public async processPromptMessage(message: PromptMessage) { return this.editorContextExtractor .extractContextForTrigger('ChatMessage') @@ -123,10 +164,8 @@ export class InlineChatProvider { const tabID = triggerEvent.tabID - const credentialsState = await AuthUtil.instance.getChatAuthState() - if ( - !(credentialsState.codewhispererChat === 'connected' && credentialsState.codewhispererCore === 'connected') - ) { + const credentialsState = AuthUtil.instance.getAuthState() + if (credentialsState !== 'connected') { const { message } = extractAuthFollowUp(credentialsState) this.errorEmitter.fire() throw new ToolkitError(message) diff --git a/packages/amazonq/src/lsp/activation.ts b/packages/amazonq/src/lsp/activation.ts index 84bae8a01a6..aebb4a60479 100644 --- a/packages/amazonq/src/lsp/activation.ts +++ b/packages/amazonq/src/lsp/activation.ts @@ -8,11 +8,11 @@ import { startLanguageServer } from './client' import { AmazonQLspInstaller } from './lspInstaller' import { lspSetupStage, ToolkitError, messages } from 'aws-core-vscode/shared' -export async function activate(ctx: vscode.ExtensionContext): Promise { +export async function activate(ctx: vscode.ExtensionContext) { try { await lspSetupStage('all', async () => { const installResult = await new AmazonQLspInstaller().resolve() - await lspSetupStage('launch', async () => await startLanguageServer(ctx, installResult.resourcePaths)) + return await lspSetupStage('launch', () => startLanguageServer(ctx, installResult.resourcePaths)) }) } catch (err) { const e = err as ToolkitError diff --git a/packages/amazonq/src/lsp/auth.ts b/packages/amazonq/src/lsp/auth.ts deleted file mode 100644 index d81f464d6a3..00000000000 --- a/packages/amazonq/src/lsp/auth.ts +++ /dev/null @@ -1,128 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - bearerCredentialsUpdateRequestType, - ConnectionMetadata, - NotificationType, - RequestType, - ResponseMessage, - UpdateCredentialsParams, -} from '@aws/language-server-runtimes/protocol' -import * as jose from 'jose' -import * as crypto from 'crypto' -import { LanguageClient } from 'vscode-languageclient' -import { AuthUtil } from 'aws-core-vscode/codewhisperer' -import { Writable } from 'stream' -import { onceChanged } from 'aws-core-vscode/utils' -import { getLogger, oneMinute } from 'aws-core-vscode/shared' -import { isSsoConnection } from 'aws-core-vscode/auth' - -export const encryptionKey = crypto.randomBytes(32) - -/** - * Sends a json payload to the language server, who is waiting to know what the encryption key is. - * Code reference: https://github.com/aws/language-servers/blob/7da212185a5da75a72ce49a1a7982983f438651a/client/vscode/src/credentialsActivation.ts#L77 - */ -export function writeEncryptionInit(stream: Writable): void { - const request = { - version: '1.0', - mode: 'JWT', - key: encryptionKey.toString('base64'), - } - stream.write(JSON.stringify(request)) - stream.write('\n') -} - -/** - * Request for custom notifications that Update Credentials and tokens. - * See core\aws-lsp-core\src\credentials\updateCredentialsRequest.ts for details - */ -export interface UpdateCredentialsRequest { - /** - * Encrypted token (JWT or PASETO) - * The token's contents differ whether IAM or Bearer token is sent - */ - data: string - /** - * Used by the runtime based language servers. - * Signals that this client will encrypt its credentials payloads. - */ - encrypted: boolean -} - -export const notificationTypes = { - updateBearerToken: new RequestType( - 'aws/credentials/token/update' - ), - deleteBearerToken: new NotificationType('aws/credentials/token/delete'), - getConnectionMetadata: new RequestType( - 'aws/credentials/getConnectionMetadata' - ), -} - -/** - * Facade over our VSCode Auth that does crud operations on the language server auth - */ -export class AmazonQLspAuth { - #logErrorIfChanged = onceChanged((s) => getLogger('amazonqLsp').error(s)) - constructor( - private readonly client: LanguageClient, - private readonly authUtil: AuthUtil = AuthUtil.instance - ) {} - - /** - * @param force bypass memoization, and forcefully update the bearer token - */ - async refreshConnection(force: boolean = false) { - const activeConnection = this.authUtil.conn - if (this.authUtil.isConnectionValid() && isSsoConnection(activeConnection)) { - // send the token to the language server - const token = await this.authUtil.getBearerToken() - await (force ? this._updateBearerToken(token) : this.updateBearerToken(token)) - } - } - - async logRefreshError(e: unknown) { - const err = e as Error - this.#logErrorIfChanged(`Unable to update bearer token: ${err.name}:${err.message}`) - } - - public updateBearerToken = onceChanged(this._updateBearerToken.bind(this)) - private async _updateBearerToken(token: string) { - const request = await this.createUpdateCredentialsRequest({ - token, - }) - - await this.client.sendRequest(bearerCredentialsUpdateRequestType.method, request) - - this.client.info(`UpdateBearerToken: ${JSON.stringify(request)}`) - } - - public startTokenRefreshInterval(pollingTime: number = oneMinute / 2) { - const interval = setInterval(async () => { - await this.refreshConnection().catch((e) => this.logRefreshError(e)) - }, pollingTime) - return interval - } - - private async createUpdateCredentialsRequest(data: any): Promise { - const payload = new TextEncoder().encode(JSON.stringify({ data })) - - const jwt = await new jose.CompactEncrypt(payload) - .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' }) - .encrypt(encryptionKey) - - return { - data: jwt, - metadata: { - sso: { - startUrl: AuthUtil.instance.startUrl, - }, - }, - encrypted: true, - } - } -} diff --git a/packages/amazonq/src/lsp/chat/activation.ts b/packages/amazonq/src/lsp/chat/activation.ts index f8e3ee16251..f6277068734 100644 --- a/packages/amazonq/src/lsp/chat/activation.ts +++ b/packages/amazonq/src/lsp/chat/activation.ts @@ -11,12 +11,20 @@ import { registerLanguageServerEventListener, registerMessageListeners } from '. import { Commands, getLogger, globals, undefinedIfEmpty } from 'aws-core-vscode/shared' import { activate as registerLegacyChatListeners } from '../../app/chat/activation' import { DefaultAmazonQAppInitContext } from 'aws-core-vscode/amazonq' -import { AuthUtil, getSelectedCustomization } from 'aws-core-vscode/codewhisperer' +import { AuthUtil, getSelectedCustomization, notifyNewCustomizations } from 'aws-core-vscode/codewhisperer' import { pushConfigUpdate } from '../config' export async function activate(languageClient: LanguageClient, encryptionKey: Buffer, mynahUIPath: string) { const disposables = globals.context.subscriptions + // Make sure we've sent an auth profile to the language server before even initializing the UI + await pushConfigUpdate(languageClient, { + type: 'profile', + profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, + }) + + await initializeCustomizations() + const provider = new AmazonQChatViewProvider(mynahUIPath) disposables.push( @@ -68,10 +76,7 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu await provider.refreshWebview() }), Commands.register('aws.amazonq.updateCustomizations', () => { - void pushConfigUpdate(languageClient, { - type: 'customization', - customization: undefinedIfEmpty(getSelectedCustomization().arn), - }) + pushCustomizationToServer(languageClient) }), globals.logOutputChannel.onDidChangeLogLevel((logLevel) => { getLogger('amazonqLsp').info(`Local log level changed to ${logLevel}, notifying LSP`) @@ -80,4 +85,35 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu }) }) ) + + /** + * Initialize customizations on extension startup + */ + async function initializeCustomizations() { + /** + * Even though this function is called "notify", it has a side effect that first restores the + * cached customization. So for {@link getSelectedCustomization()} to work as expected, we must + * call {@link notifyNewCustomizations} first. + * + * TODO: Separate restoring and notifying, or just rename the function to something better + */ + if (AuthUtil.instance.isIdcConnection() && AuthUtil.instance.isConnected()) { + await notifyNewCustomizations() + } + + /** + * HACK: We must explicitly push the customization here since restoring the customization from cache + * does not currently trigger a push to server. + * + * TODO: Always push to server whenever restoring from cache. + */ + pushCustomizationToServer(languageClient) + } + + function pushCustomizationToServer(languageClient: LanguageClient) { + void pushConfigUpdate(languageClient, { + type: 'customization', + customization: undefinedIfEmpty(getSelectedCustomization().arn), + }) + } } diff --git a/packages/amazonq/src/lsp/chat/messages.ts b/packages/amazonq/src/lsp/chat/messages.ts index 89d221e9442..11f00be3a22 100644 --- a/packages/amazonq/src/lsp/chat/messages.ts +++ b/packages/amazonq/src/lsp/chat/messages.ts @@ -38,9 +38,6 @@ import { ShowSaveFileDialogParams, LSPErrorCodes, tabBarActionRequestType, - ShowDocumentParams, - ShowDocumentResult, - ShowDocumentRequest, contextCommandsNotificationType, ContextCommandParams, openFileDiffNotificationType, @@ -57,9 +54,8 @@ import { import { v4 as uuidv4 } from 'uuid' import * as vscode from 'vscode' import { Disposable, LanguageClient, Position, TextDocumentIdentifier } from 'vscode-languageclient' -import * as jose from 'jose' import { AmazonQChatViewProvider } from './webviewProvider' -import { AuthUtil, ReferenceLogViewProvider } from 'aws-core-vscode/codewhisperer' +import { AuthUtil, CodeWhispererSettings, ReferenceLogViewProvider } from 'aws-core-vscode/codewhisperer' import { amazonQDiffScheme, AmazonQPromptSettings, messages, openUrl } from 'aws-core-vscode/shared' import { DefaultAmazonQAppInitContext, @@ -70,6 +66,8 @@ import { } from 'aws-core-vscode/amazonq' import { telemetry, TelemetryBase } from 'aws-core-vscode/telemetry' import { isValidResponseError } from './error' +import { decryptResponse, encryptRequest } from '../encryption' +import { getCursorState } from '../utils' export function registerLanguageServerEventListener(languageClient: LanguageClient, provider: AmazonQChatViewProvider) { languageClient.info( @@ -96,27 +94,20 @@ export function registerLanguageServerEventListener(languageClient: LanguageClie const telemetryName: string = e.name if (telemetryName in telemetry) { + switch (telemetryName) { + case 'codewhisperer_serviceInvocation': { + // this feature is entirely client side right now + e.data.codewhispererImportRecommendationEnabled = + CodeWhispererSettings.instance.isImportRecommendationEnabled() + break + } + } languageClient.info(`[Telemetry] Emitting ${telemetryName} telemetry: ${JSON.stringify(e.data)}`) telemetry[telemetryName as keyof TelemetryBase].emit(e.data) } }) } -function getCursorState(selection: readonly vscode.Selection[]) { - return selection.map((s) => ({ - range: { - start: { - line: s.start.line, - character: s.start.character, - }, - end: { - line: s.end.line, - character: s.end.character, - }, - }, - })) -} - export function registerMessageListeners( languageClient: LanguageClient, provider: AmazonQChatViewProvider, @@ -179,7 +170,7 @@ export function registerMessageListeners( if (fullAuthTypes.includes(authType)) { try { - await AuthUtil.instance.secondaryAuth.deleteConnection() + await AuthUtil.instance.logout() } catch (e) { languageClient.error( `[VSCode Client] Failed to authenticate after AUTH_FOLLOW_UP_CLICKED: ${(e as Error).message}` @@ -222,21 +213,12 @@ export function registerMessageListeners( const cancellationToken = new CancellationTokenSource() chatStreamTokens.set(chatParams.tabId, cancellationToken) - const chatDisposable = languageClient.onProgress( - chatRequestType, - partialResultToken, - (partialResult) => { - // Store the latest partial result - if (typeof partialResult === 'string' && encryptionKey) { - void decodeRequest(partialResult, encryptionKey).then( - (decoded) => (lastPartialResult = decoded) - ) - } else { - lastPartialResult = partialResult as ChatResult + const chatDisposable = languageClient.onProgress(chatRequestType, partialResultToken, (partialResult) => + handlePartialResult(partialResult, encryptionKey, provider, chatParams.tabId).then( + (result) => { + lastPartialResult = result } - - void handlePartialResult(partialResult, encryptionKey, provider, chatParams.tabId) - } + ) ) const editor = @@ -428,23 +410,6 @@ export function registerMessageListeners( } }) - languageClient.onRequest( - ShowDocumentRequest.method, - async (params: ShowDocumentParams): Promise> => { - try { - const uri = vscode.Uri.parse(params.uri) - const doc = await vscode.workspace.openTextDocument(uri) - await vscode.window.showTextDocument(doc, { preview: false }) - return params - } catch (e) { - return new ResponseError( - LSPErrorCodes.RequestFailed, - `Failed to open document: ${(e as Error).message}` - ) - } - } - ) - languageClient.onNotification(contextCommandsNotificationType.method, (params: ContextCommandParams) => { void provider.webview?.postMessage({ command: contextCommandsNotificationType.method, @@ -492,29 +457,6 @@ function isServerEvent(command: string) { return command.startsWith('aws/chat/') || command === 'telemetry/event' } -async function encryptRequest(params: T, encryptionKey: Buffer): Promise<{ message: string } | T> { - const payload = new TextEncoder().encode(JSON.stringify(params)) - - const encryptedMessage = await new jose.CompactEncrypt(payload) - .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' }) - .encrypt(encryptionKey) - - return { message: encryptedMessage } -} - -async function decodeRequest(request: string, key: Buffer): Promise { - const result = await jose.jwtDecrypt(request, key, { - clockTolerance: 60, // Allow up to 60 seconds to account for clock differences - contentEncryptionAlgorithms: ['A256GCM'], - keyManagementAlgorithms: ['dir'], - }) - - if (!result.payload) { - throw new Error('JWT payload not found') - } - return result.payload as T -} - /** * Decodes partial chat responses from the language server before sending them to mynah UI */ @@ -524,10 +466,7 @@ async function handlePartialResult( provider: AmazonQChatViewProvider, tabId: string ) { - const decryptedMessage = - typeof partialResult === 'string' && encryptionKey - ? await decodeRequest(partialResult, encryptionKey) - : (partialResult as T) + const decryptedMessage = await decryptResponse(partialResult, encryptionKey) if (decryptedMessage.body !== undefined) { void provider.webview?.postMessage({ @@ -537,6 +476,7 @@ async function handlePartialResult( tabId: tabId, }) } + return decryptedMessage } /** @@ -550,8 +490,8 @@ async function handleCompleteResult( tabId: string, disposable: Disposable ) { - const decryptedMessage = - typeof result === 'string' && encryptionKey ? await decodeRequest(result, encryptionKey) : (result as T) + const decryptedMessage = await decryptResponse(result, encryptionKey) + void provider.webview?.postMessage({ command: chatRequestType.method, params: decryptedMessage, diff --git a/packages/amazonq/src/lsp/chat/webviewProvider.ts b/packages/amazonq/src/lsp/chat/webviewProvider.ts index 1a513f1df3f..3f9d273ab96 100644 --- a/packages/amazonq/src/lsp/chat/webviewProvider.ts +++ b/packages/amazonq/src/lsp/chat/webviewProvider.ts @@ -137,7 +137,7 @@ export class AmazonQChatViewProvider implements WebviewViewProvider { let qChat = undefined const init = () => { const vscodeApi = acquireVsCodeApi() - const hybridChatConnector = new HybridChatAdapter(${(await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected'},${featureConfigData},${welcomeCount},${disclaimerAcknowledged},${regionProfileString},${disabledCommands},${isSMUS},${isSM},vscodeApi.postMessage) + const hybridChatConnector = new HybridChatAdapter(${AuthUtil.instance.isConnected()},${featureConfigData},${welcomeCount},${disclaimerAcknowledged},${regionProfileString},${disabledCommands},${isSMUS},${isSM},vscodeApi.postMessage) const commands = [hybridChatConnector.initialQuickActions[0]] qChat = amazonQChat.createChat(vscodeApi, {disclaimerAcknowledged: ${disclaimerAcknowledged}, pairProgrammingAcknowledged: ${pairProgrammingAcknowledged}, agenticMode: true, quickActionCommands: commands}, hybridChatConnector, ${JSON.stringify(featureConfigData)}); } diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index 549b0ac7dad..08d049c9aba 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -5,9 +5,9 @@ import vscode, { env, version } from 'vscode' import * as nls from 'vscode-nls' +import * as crypto from 'crypto' +import * as jose from 'jose' import { LanguageClient, LanguageClientOptions, RequestType, State } from 'vscode-languageclient' -import { InlineCompletionManager } from '../app/inline/completion' -import { AmazonQLspAuth, encryptionKey, notificationTypes } from './auth' import { CreateFilesParams, DeleteFilesParams, @@ -17,43 +17,72 @@ import { RenameFilesParams, ResponseMessage, WorkspaceFolder, + GetSsoTokenProgress, + GetSsoTokenProgressToken, + GetSsoTokenProgressType, + MessageActionItem, + ShowMessageRequest, + ShowMessageRequestParams, + ConnectionMetadata, + ShowDocumentRequest, + ShowDocumentParams, + ShowDocumentResult, + ResponseError, + LSPErrorCodes, + updateConfigurationRequestType, } from '@aws/language-server-runtimes/protocol' -import { AuthUtil, CodeWhispererSettings, getSelectedCustomization } from 'aws-core-vscode/codewhisperer' +import { + AuthUtil, + CodeWhispererSettings, + getSelectedCustomization, + TelemetryHelper, +} from 'aws-core-vscode/codewhisperer' import { Settings, createServerOptions, globals, Experiments, - Commands, - oneSecond, validateNodeExe, getLogger, undefinedIfEmpty, getOptOutPreference, isAmazonLinux2, + oidcClientName, + openUrl, getClientId, extensionVersion, + Commands, } from 'aws-core-vscode/shared' import { processUtils } from 'aws-core-vscode/shared' -import { activate } from './chat/activation' +import { activate as activateChat } from './chat/activation' +import { activate as activeInlineChat } from '../inlineChat/activation' import { AmazonQResourcePaths } from './lspInstaller' +import { auth2 } from 'aws-core-vscode/auth' import { ConfigSection, isValidConfigSection, pushConfigUpdate, toAmazonQLSPLogLevel } from './config' import { telemetry } from 'aws-core-vscode/telemetry' +import { SessionManager } from '../app/inline/sessionManager' +import { LineTracker } from '../app/inline/stateTracker/lineTracker' +import { InlineChatTutorialAnnotation } from '../app/inline/tutorials/inlineChatTutorialAnnotation' +import { InlineTutorialAnnotation } from '../app/inline/tutorials/inlineTutorialAnnotation' +import { InlineCompletionManager } from '../app/inline/completion' const localize = nls.loadMessageBundle() const logger = getLogger('amazonqLsp.lspClient') export const glibcLinker: string = process.env.VSCODE_SERVER_CUSTOM_GLIBC_LINKER || '' export const glibcPath: string = process.env.VSCODE_SERVER_CUSTOM_GLIBC_PATH || '' - export function hasGlibcPatch(): boolean { return glibcLinker.length > 0 && glibcPath.length > 0 } +export const clientId = 'amazonq' +export const clientName = oidcClientName() +export const encryptionKey = crypto.randomBytes(32) + export async function startLanguageServer( extensionContext: vscode.ExtensionContext, resourcePaths: AmazonQResourcePaths -) { +): Promise { const toDispose = extensionContext.subscriptions const serverModule = resourcePaths.lsp @@ -67,8 +96,6 @@ export async function startLanguageServer( ] const documentSelector = [{ scheme: 'file', language: '*' }] - - const clientId = 'amazonq' const traceServerEnabled = Settings.instance.isSet(`${clientId}.trace.server`) let executable: string[] = [] // apply the GLIBC 2.28 path to node js runtime binary @@ -91,6 +118,7 @@ export async function startLanguageServer( await validateNodeExe(executable, resourcePaths.lsp, argv, logger) // Options to control the language client + const clientName = 'AmazonQ-For-VSCode' const clientOptions: LanguageClientOptions = { // Register the server for json documents documentSelector, @@ -115,7 +143,7 @@ export async function startLanguageServer( name: env.appName, version: version, extension: { - name: 'AmazonQ-For-VSCode', + name: clientName, version: extensionVersion, }, clientId: getClientId(globals.globalState), @@ -150,76 +178,209 @@ export async function startLanguageServer( }), } - const client = new LanguageClient( - clientId, - localize('amazonq.server.name', 'Amazon Q Language Server'), - serverOptions, - clientOptions - ) + const lspName = localize('amazonq.server.name', 'Amazon Q Language Server') + const client = new LanguageClient(clientId, lspName, serverOptions, clientOptions) const disposable = client.start() toDispose.push(disposable) + await client.onReady() - const auth = await initializeAuth(client) + /** + * We use the Flare Auth language server, and our Auth client depends on it. + * Because of this we initialize our Auth client **immediately** after the language server is ready. + * Doing this removes the chance of something else attempting to use the Auth client before it is ready. + */ + await initializeAuth(client) - await onLanguageServerReady(auth, client, resourcePaths, toDispose) + await postStartLanguageServer(extensionContext, client, resourcePaths, toDispose) return client + + async function initializeAuth(client: LanguageClient) { + AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey)) + + // Migrate SSO connections from old Auth to the LSP identity server + // This function only migrates connections once + // This call can be removed once all/most users have updated to the latest AmazonQ version + try { + await AuthUtil.instance.migrateSsoConnectionToLsp(clientName) + } catch (e) { + getLogger().error(`Error while migration SSO connection to Amazon Q LSP: ${e}`) + } + + /** All must be setup before {@link AuthUtil.restore} otherwise they may not trigger when expected */ + AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => { + void pushConfigUpdate(client, { + type: 'profile', + profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, + }) + }) + + // Try and restore a cached connection if exists + await AuthUtil.instance.restore() + } } -async function initializeAuth(client: LanguageClient): Promise { - const auth = new AmazonQLspAuth(client) - await auth.refreshConnection(true) - return auth +async function setupInline( + extensionContext: vscode.ExtensionContext, + client: LanguageClient, + toDispose: vscode.Disposable[] +) { + const sessionManager = new SessionManager() + const lineTracker = new LineTracker() + const inlineTutorialAnnotation = new InlineTutorialAnnotation(lineTracker, sessionManager) + const inlineChatTutorialAnnotation = new InlineChatTutorialAnnotation(inlineTutorialAnnotation) + + const inlineManager = new InlineCompletionManager(client, sessionManager, lineTracker, inlineTutorialAnnotation) + + inlineManager.registerInlineCompletion() + + activeInlineChat(extensionContext, client, encryptionKey, inlineChatTutorialAnnotation) + + toDispose.push( + inlineManager, + Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => { + await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger') + }), + Commands.register('aws.amazonq.refreshAnnotation', async (forceProceed: boolean) => { + telemetry.record({ + traceId: TelemetryHelper.instance.traceId, + }) + + const editor = vscode.window.activeTextEditor + if (editor) { + if (forceProceed) { + await inlineTutorialAnnotation.refresh(editor, 'codewhisperer', true) + } else { + await inlineTutorialAnnotation.refresh(editor, 'codewhisperer') + } + } + }), + Commands.register('aws.amazonq.dismissTutorial', async () => { + const editor = vscode.window.activeTextEditor + if (editor) { + inlineTutorialAnnotation.clear() + try { + telemetry.ui_click.emit({ elementId: `dismiss_${inlineTutorialAnnotation.currentState.id}` }) + } catch (_) {} + await inlineTutorialAnnotation.dismissTutorial() + getLogger().debug(`codewhisperer: user dismiss tutorial.`) + } + }), + vscode.workspace.onDidCloseTextDocument(async () => { + await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion') + }) + ) } -async function onLanguageServerReady( - auth: AmazonQLspAuth, +async function postStartLanguageServer( + extensionContext: vscode.ExtensionContext, client: LanguageClient, resourcePaths: AmazonQResourcePaths, toDispose: vscode.Disposable[] ) { - if (Experiments.instance.get('amazonqLSPInline', false)) { - const inlineManager = new InlineCompletionManager(client) - inlineManager.registerInlineCompletion() - toDispose.push( - inlineManager, - Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => { - await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger') - }), - vscode.workspace.onDidCloseTextDocument(async () => { - await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion') + // Request handler for when the server wants to know about the clients auth connnection. Must be registered before the initial auth init call + client.onRequest(auth2.notificationTypes.getConnectionMetadata.method, () => { + return { + sso: { + startUrl: AuthUtil.instance.connection?.startUrl, + }, + } + }) + + client.onRequest( + ShowMessageRequest.method, + async (params: ShowMessageRequestParams) => { + const actions = params.actions?.map((a) => a.title) ?? [] + const response = await vscode.window.showInformationMessage(params.message, { modal: true }, ...actions) + return params.actions?.find((a) => a.title === response) ?? (undefined as unknown as null) + } + ) + + client.onRequest( + ShowDocumentRequest.method, + async (params: ShowDocumentParams): Promise> => { + const uri = vscode.Uri.parse(params.uri) + getLogger().info(`Processing ShowDocumentRequest for URI scheme: ${uri.scheme}`) + try { + if (uri.scheme.startsWith('http')) { + getLogger().info('Opening URL...') + await openUrl(vscode.Uri.parse(params.uri)) + } else { + getLogger().info('Opening text document...') + const doc = await vscode.workspace.openTextDocument(uri) + await vscode.window.showTextDocument(doc, { preview: false }) + } + return params + } catch (e) { + return new ResponseError( + LSPErrorCodes.RequestFailed, + `Failed to process ShowDocumentRequest: ${(e as Error).message}` + ) + } + } + ) + + const sendProfileToLsp = async () => { + try { + const result = await client.sendRequest(updateConfigurationRequestType.method, { + section: 'aws.q', + settings: { + profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, + }, }) - ) + client.info( + `Client: Updated Amazon Q Profile ${AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn} to Amazon Q LSP`, + result + ) + } catch (err) { + client.error('Error when setting Q Developer Profile to Amazon Q LSP', err) + } } - if (Experiments.instance.get('amazonqChatLSP', true)) { - await activate(client, encryptionKey, resourcePaths.ui) - } + let promise: Promise | undefined + let resolver: () => void = () => {} + client.onProgress(GetSsoTokenProgressType, GetSsoTokenProgressToken, async (partialResult: GetSsoTokenProgress) => { + const decryptedKey = await jose.compactDecrypt(partialResult as unknown as string, encryptionKey) + const val: GetSsoTokenProgress = JSON.parse(decryptedKey.plaintext.toString()) - const refreshInterval = auth.startTokenRefreshInterval(10 * oneSecond) + if (val.state === 'InProgress') { + if (promise) { + resolver() + } + promise = new Promise((resolve) => { + resolver = resolve + }) + } else { + resolver() + promise = undefined + return + } - // We manually push the cached values the first time since event handlers, which should push, may not have been setup yet. - // Execution order is weird and should be fixed in the flare implementation. - // TODO: Revisit if we need this if we setup the event handlers properly - if (AuthUtil.instance.isConnectionValid()) { - await sendProfileToLsp(client) + // send profile to lsp once. + void sendProfileToLsp() - await pushConfigUpdate(client, { - type: 'customization', - customization: getSelectedCustomization(), - }) + void vscode.window.withProgress( + { + cancellable: true, + location: vscode.ProgressLocation.Notification, + title: val.message, + }, + async (_) => { + await promise + } + ) + }) + + if (Experiments.instance.get('amazonqChatLSP', true)) { + await activateChat(client, encryptionKey, resourcePaths.ui) } + await setupInline(extensionContext, client, toDispose) + toDispose.push( - AuthUtil.instance.auth.onDidChangeActiveConnection(async () => { - await auth.refreshConnection() - }), - AuthUtil.instance.auth.onDidDeleteConnection(async () => { - client.sendNotification(notificationTypes.deleteBearerToken.method) - }), - AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(() => sendProfileToLsp(client)), + AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(sendProfileToLsp), vscode.commands.registerCommand('aws.amazonq.getWorkspaceId', async () => { const requestType = new RequestType( 'aws/getConfigurationFromServer' @@ -275,24 +436,16 @@ async function onLanguageServerReady( }, } as DidChangeWorkspaceFoldersParams) }), - { dispose: () => clearInterval(refreshInterval) }, // Set this inside onReady so that it only triggers on subsequent language server starts (not the first) - onServerRestartHandler(client, auth) + onServerRestartHandler(client) ) - - async function sendProfileToLsp(client: LanguageClient) { - await pushConfigUpdate(client, { - type: 'profile', - profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, - }) - } } /** * When the server restarts (likely due to a crash, then the LanguageClient automatically starts it again) * we need to run some server intialization again. */ -function onServerRestartHandler(client: LanguageClient, auth: AmazonQLspAuth) { +function onServerRestartHandler(client: LanguageClient) { return client.onDidChangeState(async (e) => { // Ensure we are in a "restart" state if (!(e.oldState === State.Starting && e.newState === State.Running)) { @@ -306,7 +459,7 @@ function onServerRestartHandler(client: LanguageClient, auth: AmazonQLspAuth) { telemetry.languageServer_crash.emit({ id: 'AmazonQ' }) // Need to set the auth token in the again - await auth.refreshConnection(true) + await AuthUtil.instance.restore() }) } diff --git a/packages/amazonq/src/lsp/config.ts b/packages/amazonq/src/lsp/config.ts index 1760fb51401..66edc9ff6f1 100644 --- a/packages/amazonq/src/lsp/config.ts +++ b/packages/amazonq/src/lsp/config.ts @@ -3,15 +3,14 @@ * SPDX-License-Identifier: Apache-2.0 */ import * as vscode from 'vscode' -import { DevSettings, getServiceEnvVarConfig } from 'aws-core-vscode/shared' -import { LspConfig } from 'aws-core-vscode/amazonq' +import { DevSettings, getServiceEnvVarConfig, BaseLspInstaller } from 'aws-core-vscode/shared' import { LanguageClient } from 'vscode-languageclient' import { DidChangeConfigurationNotification, updateConfigurationRequestType, } from '@aws/language-server-runtimes/protocol' -export interface ExtendedAmazonQLSPConfig extends LspConfig { +export interface ExtendedAmazonQLSPConfig extends BaseLspInstaller.LspConfig { ui?: string } diff --git a/packages/amazonq/src/lsp/encryption.ts b/packages/amazonq/src/lsp/encryption.ts new file mode 100644 index 00000000000..246c64f476b --- /dev/null +++ b/packages/amazonq/src/lsp/encryption.ts @@ -0,0 +1,34 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +import * as jose from 'jose' + +export async function encryptRequest(params: T, encryptionKey: Buffer): Promise<{ message: string } | T> { + const payload = new TextEncoder().encode(JSON.stringify(params)) + + const encryptedMessage = await new jose.CompactEncrypt(payload) + .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' }) + .encrypt(encryptionKey) + + return { message: encryptedMessage } +} + +export async function decryptResponse(response: unknown, key: Buffer | undefined) { + // Note that casts are required since language client requests return 'unknown' type. + // If we can't decrypt, return original response casted. + if (typeof response !== 'string' || key === undefined) { + return response as T + } + + const result = await jose.jwtDecrypt(response, key, { + clockTolerance: 60, // Allow up to 60 seconds to account for clock differences + contentEncryptionAlgorithms: ['A256GCM'], + keyManagementAlgorithms: ['dir'], + }) + + if (!result.payload) { + throw new Error('JWT payload not found') + } + return result.payload as T +} diff --git a/packages/amazonq/src/lsp/utils.ts b/packages/amazonq/src/lsp/utils.ts new file mode 100644 index 00000000000..f5b010c536b --- /dev/null +++ b/packages/amazonq/src/lsp/utils.ts @@ -0,0 +1,26 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +import * as vscode from 'vscode' +import { CursorState } from '@aws/language-server-runtimes-types' + +/** + * Convert from vscode selection type to the general CursorState expected by the AmazonQLSP. + * @param selection + * @returns + */ +export function getCursorState(selection: readonly vscode.Selection[]): CursorState[] { + return selection.map((s) => ({ + range: { + start: { + line: s.start.line, + character: s.start.character, + }, + end: { + line: s.end.line, + character: s.end.character, + }, + }, + })) +} diff --git a/packages/amazonq/src/util/clearCache.ts b/packages/amazonq/src/util/clearCache.ts index 8c93b35ac12..c6d5f8d6839 100644 --- a/packages/amazonq/src/util/clearCache.ts +++ b/packages/amazonq/src/util/clearCache.ts @@ -32,10 +32,9 @@ async function clearCache() { return } - // SSO cache persists on disk, this should indirectly delete it - const conn = AuthUtil.instance.conn - if (conn) { - await AuthUtil.instance.auth.deleteConnection(conn) + // SSO cache persists on disk, this should log out + if (AuthUtil.instance.isConnected()) { + await AuthUtil.instance.logout() } await globals.globalState.clear() diff --git a/packages/amazonq/test/e2e/amazonq/utils/setup.ts b/packages/amazonq/test/e2e/amazonq/utils/setup.ts index dd82b1f0b19..ef7ba540198 100644 --- a/packages/amazonq/test/e2e/amazonq/utils/setup.ts +++ b/packages/amazonq/test/e2e/amazonq/utils/setup.ts @@ -5,9 +5,9 @@ import { AuthUtil } from 'aws-core-vscode/codewhisperer' export async function loginToIdC() { - const authState = await AuthUtil.instance.getChatAuthState() + const authState = AuthUtil.instance.getAuthState() if (process.env['AWS_TOOLKIT_AUTOMATION'] === 'local') { - if (authState.amazonQ !== 'connected') { + if (authState !== 'connected') { throw new Error('You will need to login manually before running tests.') } return @@ -22,5 +22,5 @@ export async function loginToIdC() { ) } - await AuthUtil.instance.connectToEnterpriseSso(startUrl, region) + await AuthUtil.instance.login(startUrl, region) } diff --git a/packages/amazonq/test/e2e/inline/inline.test.ts b/packages/amazonq/test/e2e/inline/inline.test.ts index 43a9f67ab73..bcc41851eca 100644 --- a/packages/amazonq/test/e2e/inline/inline.test.ts +++ b/packages/amazonq/test/e2e/inline/inline.test.ts @@ -5,18 +5,10 @@ import * as vscode from 'vscode' import assert from 'assert' -import { - closeAllEditors, - getTestWindow, - registerAuthHook, - resetCodeWhispererGlobalVariables, - TestFolder, - toTextEditor, - using, -} from 'aws-core-vscode/test' -import { RecommendationHandler, RecommendationService, session } from 'aws-core-vscode/codewhisperer' +import { closeAllEditors, registerAuthHook, TestFolder, toTextEditor, using } from 'aws-core-vscode/test' import { Commands, globals, sleep, waitUntil, collectionUtil } from 'aws-core-vscode/shared' import { loginToIdC } from '../amazonq/utils/setup' +import { vsCodeState } from 'aws-core-vscode/codewhisperer' describe('Amazon Q Inline', async function () { const retries = 3 @@ -40,7 +32,6 @@ describe('Amazon Q Inline', async function () { const folder = await TestFolder.create() tempFolder = folder.path await closeAllEditors() - await resetCodeWhispererGlobalVariables() }) afterEach(async function () { @@ -54,7 +45,6 @@ describe('Amazon Q Inline', async function () { const events = getUserTriggerDecision() console.table({ 'telemetry events': JSON.stringify(events), - 'recommendation service status': RecommendationService.instance.isRunning, }) } @@ -71,31 +61,6 @@ describe('Amazon Q Inline', async function () { }) } - async function waitForRecommendations() { - const suggestionShown = await waitUntil(async () => session.getSuggestionState(0) === 'Showed', waitOptions) - if (!suggestionShown) { - throw new Error(`Suggestion did not show. Suggestion States: ${JSON.stringify(session.suggestionStates)}`) - } - const suggestionVisible = await waitUntil( - async () => RecommendationHandler.instance.isSuggestionVisible(), - waitOptions - ) - if (!suggestionVisible) { - throw new Error( - `Suggestions failed to become visible. Suggestion States: ${JSON.stringify(session.suggestionStates)}` - ) - } - console.table({ - 'suggestions states': JSON.stringify(session.suggestionStates), - 'valid recommendation': RecommendationHandler.instance.isValidResponse(), - 'recommendation service status': RecommendationService.instance.isRunning, - recommendations: session.recommendations, - }) - if (!RecommendationHandler.instance.isValidResponse()) { - throw new Error('Did not find a valid response') - } - } - /** * Waits for a specific telemetry event to be emitted with the expected suggestion state. * It looks like there might be a potential race condition in codewhisperer causing telemetry @@ -149,8 +114,9 @@ describe('Amazon Q Inline', async function () { await invokeCompletion() originalEditorContents = vscode.window.activeTextEditor?.document.getText() - // wait until the ghost text appears - await waitForRecommendations() + // wait until all the recommendations have finished + await waitUntil(() => Promise.resolve(vsCodeState.isRecommendationsActive === true), waitOptions) + await waitUntil(() => Promise.resolve(vsCodeState.isRecommendationsActive === false), waitOptions) } beforeEach(async () => { @@ -163,14 +129,12 @@ describe('Amazon Q Inline', async function () { try { await setup() console.log(`test run ${attempt} succeeded`) - logUserDecisionStatus() break } catch (e) { console.log(`test run ${attempt} failed`) console.log(e) logUserDecisionStatus() attempt++ - await resetCodeWhispererGlobalVariables() } } if (attempt === retries) { @@ -216,29 +180,6 @@ describe('Amazon Q Inline', async function () { assert.deepStrictEqual(vscode.window.activeTextEditor?.document.getText(), originalEditorContents) }) }) - - it(`${name} invoke on unsupported filetype`, async function () { - await setupEditor({ - name: 'test.zig', - contents: `fn doSomething() void { - - }`, - }) - - /** - * Add delay between editor loading and invoking completion - * @see beforeEach in supported filetypes for more information - */ - await sleep(1000) - await invokeCompletion() - - if (name === 'automatic') { - // It should never get triggered since its not a supported file type - assert.deepStrictEqual(RecommendationService.instance.isRunning, false) - } else { - await getTestWindow().waitForMessage('currently not supported by Amazon Q inline suggestions') - } - }) }) } }) diff --git a/packages/amazonq/test/e2e/lsp/amazonqLsp.test.ts b/packages/amazonq/test/e2e/lsp/amazonqLsp.test.ts index d3e90ec4e8e..f4a60ff282b 100644 --- a/packages/amazonq/test/e2e/lsp/amazonqLsp.test.ts +++ b/packages/amazonq/test/e2e/lsp/amazonqLsp.test.ts @@ -6,13 +6,13 @@ import { AmazonQLspInstaller } from '../../../src/lsp/lspInstaller' import { defaultAmazonQLspConfig } from '../../../src/lsp/config' import { createLspInstallerTests } from './lspInstallerUtil' -import { LspConfig } from 'aws-core-vscode/amazonq' +import { BaseLspInstaller } from 'aws-core-vscode/shared' describe('AmazonQLSP', () => { createLspInstallerTests({ suiteName: 'AmazonQLSPInstaller', lspConfig: defaultAmazonQLspConfig, - createInstaller: (lspConfig?: LspConfig) => new AmazonQLspInstaller(lspConfig), + createInstaller: (lspConfig?: BaseLspInstaller.LspConfig) => new AmazonQLspInstaller(lspConfig), targetContents: [ { bytes: 0, diff --git a/packages/amazonq/test/e2e/lsp/lspInstallerUtil.ts b/packages/amazonq/test/e2e/lsp/lspInstallerUtil.ts index c7ca7a4ff9b..d4251959756 100644 --- a/packages/amazonq/test/e2e/lsp/lspInstallerUtil.ts +++ b/packages/amazonq/test/e2e/lsp/lspInstallerUtil.ts @@ -18,7 +18,6 @@ import { } from 'aws-core-vscode/shared' import * as semver from 'semver' import { assertTelemetry } from 'aws-core-vscode/test' -import { LspConfig, LspController } from 'aws-core-vscode/amazonq' import { LanguageServerSetup } from 'aws-core-vscode/telemetry' function createVersion(version: string, contents: TargetContent[]) { @@ -44,8 +43,8 @@ export function createLspInstallerTests({ resetEnv, }: { suiteName: string - lspConfig: LspConfig - createInstaller: (lspConfig?: LspConfig) => BaseLspInstaller.BaseLspInstaller + lspConfig: BaseLspInstaller.LspConfig + createInstaller: (lspConfig?: BaseLspInstaller.LspConfig) => BaseLspInstaller.BaseLspInstaller targetContents: TargetContent[] setEnv: (path: string) => void resetEnv: () => void @@ -60,8 +59,6 @@ export function createLspInstallerTests({ installer = createInstaller() tempDir = await makeTemporaryToolkitFolder() sandbox.stub(LanguageServerResolver.prototype, 'defaultDownloadFolder').returns(tempDir) - // Called on extension activation and can contaminate telemetry. - sandbox.stub(LspController.prototype, 'trySetupLsp') }) afterEach(async () => { diff --git a/packages/amazonq/test/e2e/lsp/workspaceContextLsp.test.ts b/packages/amazonq/test/e2e/lsp/workspaceContextLsp.test.ts deleted file mode 100644 index 75d57949c0b..00000000000 --- a/packages/amazonq/test/e2e/lsp/workspaceContextLsp.test.ts +++ /dev/null @@ -1,42 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as os from 'os' -import { createLspInstallerTests } from './lspInstallerUtil' -import { defaultAmazonQWorkspaceLspConfig, LspClient, LspConfig, WorkspaceLspInstaller } from 'aws-core-vscode/amazonq' -import assert from 'assert' - -describe('AmazonQWorkspaceLSP', () => { - createLspInstallerTests({ - suiteName: 'AmazonQWorkspaceLSPInstaller', - lspConfig: defaultAmazonQWorkspaceLspConfig, - createInstaller: (lspConfig?: LspConfig) => new WorkspaceLspInstaller.WorkspaceLspInstaller(lspConfig), - targetContents: [ - { - bytes: 0, - filename: `qserver-${os.platform()}-${os.arch()}.zip`, - hashes: [], - url: 'http://fakeurl', - }, - ], - setEnv: (path: string) => { - process.env.__AMAZONQWORKSPACELSP_PATH = path - }, - resetEnv: () => { - delete process.env.__AMAZONQWORKSPACELSP_PATH - }, - }) - - it('activates', async () => { - const ok = await LspClient.instance.waitUntilReady() - if (!ok) { - assert.fail('Workspace context language server failed to become ready') - } - const serverUsage = await LspClient.instance.getLspServerUsage() - if (!serverUsage) { - assert.fail('Unable to verify that the workspace context language server has been activated') - } - }) -}) diff --git a/packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts b/packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts index d2182329e45..a8bc854c97f 100644 --- a/packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts +++ b/packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts @@ -3,18 +3,27 @@ * SPDX-License-Identifier: Apache-2.0 */ import sinon from 'sinon' -import { CancellationToken, commands, languages, Position } from 'vscode' +import { + CancellationToken, + commands, + InlineCompletionItem, + languages, + Position, + window, + Range, + InlineCompletionTriggerKind, +} from 'vscode' import assert from 'assert' import { LanguageClient } from 'vscode-languageclient' +import { StringValue } from 'vscode-languageserver-types' import { AmazonQInlineCompletionItemProvider, InlineCompletionManager } from '../../../../../src/app/inline/completion' import { RecommendationService } from '../../../../../src/app/inline/recommendationService' import { SessionManager } from '../../../../../src/app/inline/sessionManager' -import { createMockDocument, createMockTextEditor } from 'aws-core-vscode/test' -import { - ReferenceHoverProvider, - ReferenceInlineProvider, - ReferenceLogViewProvider, -} from 'aws-core-vscode/codewhisperer' +import { createMockDocument, createMockTextEditor, getTestWindow, installFakeClock } from 'aws-core-vscode/test' +import { noInlineSuggestionsMsg, ReferenceHoverProvider, ReferenceLogViewProvider } from 'aws-core-vscode/codewhisperer' +import { InlineGeneratingMessage } from '../../../../../src/app/inline/inlineGeneratingMessage' +import { LineTracker } from '../../../../../src/app/inline/stateTracker/lineTracker' +import { InlineTutorialAnnotation } from '../../../../../src/app/inline/tutorials/inlineTutorialAnnotation' describe('InlineCompletionManager', () => { let manager: InlineCompletionManager @@ -72,7 +81,10 @@ describe('InlineCompletionManager', () => { sendNotification: sendNotificationStub, } as unknown as LanguageClient - manager = new InlineCompletionManager(languageClient) + const sessionManager = new SessionManager() + const lineTracker = new LineTracker() + const inlineTutorialAnnotation = new InlineTutorialAnnotation(lineTracker, sessionManager) + manager = new InlineCompletionManager(languageClient, sessionManager, lineTracker, inlineTutorialAnnotation) getActiveSessionStub = sandbox.stub(manager['sessionManager'], 'getActiveSession') getActiveRecommendationStub = sandbox.stub(manager['sessionManager'], 'getActiveRecommendation') getReferenceStub = sandbox.stub(ReferenceLogViewProvider, 'getReferenceLog') @@ -213,46 +225,6 @@ describe('InlineCompletionManager', () => { assert(registerProviderStub.calledTwice) // Once in constructor, once after rejection }) }) - - describe('previous command', () => { - it('should register and handle previous command correctly', async () => { - const prevCommandCall = registerCommandStub - .getCalls() - .find((call) => call.args[0] === 'editor.action.inlineSuggest.showPrevious') - - assert(prevCommandCall, 'Previous command should be registered') - - if (prevCommandCall) { - const handler = prevCommandCall.args[1] - await handler() - - assert(executeCommandStub.calledWith('editor.action.inlineSuggest.hide')) - assert(disposableStub.calledOnce) - assert(registerProviderStub.calledTwice) - assert(executeCommandStub.calledWith('editor.action.inlineSuggest.trigger')) - } - }) - }) - - describe('next command', () => { - it('should register and handle next command correctly', async () => { - const nextCommandCall = registerCommandStub - .getCalls() - .find((call) => call.args[0] === 'editor.action.inlineSuggest.showNext') - - assert(nextCommandCall, 'Next command should be registered') - - if (nextCommandCall) { - const handler = nextCommandCall.args[1] - await handler() - - assert(executeCommandStub.calledWith('editor.action.inlineSuggest.hide')) - assert(disposableStub.calledOnce) - assert(registerProviderStub.calledTwice) - assert(executeCommandStub.calledWith('editor.action.inlineSuggest.trigger')) - } - }) - }) }) describe('AmazonQInlineCompletionItemProvider', () => { @@ -261,11 +233,13 @@ describe('InlineCompletionManager', () => { let provider: AmazonQInlineCompletionItemProvider let getAllRecommendationsStub: sinon.SinonStub let recommendationService: RecommendationService - let setInlineReferenceStub: sinon.SinonStub + let inlineTutorialAnnotation: InlineTutorialAnnotation beforeEach(() => { - recommendationService = new RecommendationService(mockSessionManager) - setInlineReferenceStub = sandbox.stub(ReferenceInlineProvider.instance, 'setInlineReference') + const lineTracker = new LineTracker() + const activeStateController = new InlineGeneratingMessage(lineTracker) + inlineTutorialAnnotation = new InlineTutorialAnnotation(lineTracker, mockSessionManager) + recommendationService = new RecommendationService(mockSessionManager, activeStateController) mockSessionManager = { getActiveSession: getActiveSessionStub, @@ -281,12 +255,14 @@ describe('InlineCompletionManager', () => { getActiveRecommendationStub.returns(mockSuggestions) getAllRecommendationsStub = sandbox.stub(recommendationService, 'getAllRecommendations') getAllRecommendationsStub.resolves() + sandbox.stub(window, 'activeTextEditor').value(createMockTextEditor()) }), it('should call recommendation service to get new suggestions for new sessions', async () => { provider = new AmazonQInlineCompletionItemProvider( languageClient, recommendationService, - mockSessionManager + mockSessionManager, + inlineTutorialAnnotation ) const items = await provider.provideInlineCompletionItems( mockDocument, @@ -297,39 +273,130 @@ describe('InlineCompletionManager', () => { assert(getAllRecommendationsStub.calledOnce) assert.deepStrictEqual(items, mockSuggestions) }), - it('should not call recommendation service for existing sessions', async () => { + it('should handle reference if there is any', async () => { provider = new AmazonQInlineCompletionItemProvider( languageClient, recommendationService, mockSessionManager, - false + inlineTutorialAnnotation ) - const items = await provider.provideInlineCompletionItems( + await provider.provideInlineCompletionItems(mockDocument, mockPosition, mockContext, mockToken) + }), + it('should add a range to the completion item when missing', async function () { + provider = new AmazonQInlineCompletionItemProvider( + languageClient, + recommendationService, + mockSessionManager, + inlineTutorialAnnotation + ) + getActiveRecommendationStub.returns([ + { + insertText: 'testText', + itemId: 'itemId', + }, + { + insertText: 'testText2', + itemId: 'itemId2', + range: undefined, + }, + ]) + const cursorPosition = new Position(5, 6) + const result = await provider.provideInlineCompletionItems( + mockDocument, + cursorPosition, + mockContext, + mockToken + ) + + for (const item of result) { + assert.deepStrictEqual(item.range, new Range(cursorPosition, cursorPosition)) + } + }), + it('should handle StringValue instead of strings', async function () { + provider = new AmazonQInlineCompletionItemProvider( + languageClient, + recommendationService, + mockSessionManager, + inlineTutorialAnnotation + ) + const expectedText = 'this is my text' + getActiveRecommendationStub.returns([ + { + insertText: { kind: 'snippet', value: 'this is my text' } satisfies StringValue, + itemId: 'itemId', + }, + ]) + const result = await provider.provideInlineCompletionItems( mockDocument, mockPosition, mockContext, mockToken ) - assert(getAllRecommendationsStub.notCalled) - assert.deepStrictEqual(items, mockSuggestions) + + assert.strictEqual(result[0].insertText, expectedText) }), - it('should handle reference if there is any', async () => { + it('shows message to user when manual invoke fails to produce results', async function () { provider = new AmazonQInlineCompletionItemProvider( languageClient, recommendationService, mockSessionManager, - false + inlineTutorialAnnotation ) - await provider.provideInlineCompletionItems(mockDocument, mockPosition, mockContext, mockToken) - assert(setInlineReferenceStub.calledOnce) - assert( - setInlineReferenceStub.calledWithExactly( - mockPosition.line, - mockSuggestions[0].insertText, - fakeReferences - ) + getActiveRecommendationStub.returns([]) + const messageShown = new Promise((resolve) => + getTestWindow().onDidShowMessage((e) => { + assert.strictEqual(e.message, noInlineSuggestionsMsg) + resolve(true) + }) + ) + await provider.provideInlineCompletionItems( + mockDocument, + mockPosition, + { triggerKind: InlineCompletionTriggerKind.Invoke, selectedCompletionInfo: undefined }, + mockToken ) + await messageShown + }) + describe('debounce behavior', function () { + let clock: ReturnType + + beforeEach(function () { + clock = installFakeClock() }) + + after(function () { + clock.uninstall() + }) + + it('should only trigger once on rapid events', async () => { + provider = new AmazonQInlineCompletionItemProvider( + languageClient, + recommendationService, + mockSessionManager, + inlineTutorialAnnotation + ) + const p1 = provider.provideInlineCompletionItems(mockDocument, mockPosition, mockContext, mockToken) + const p2 = provider.provideInlineCompletionItems(mockDocument, mockPosition, mockContext, mockToken) + const p3 = provider.provideInlineCompletionItems( + mockDocument, + new Position(2, 2), + mockContext, + mockToken + ) + + await clock.tickAsync(1000) + + // All promises should be the same object when debounced properly. + assert.strictEqual(p1, p2) + assert.strictEqual(p1, p3) + await p1 + await p2 + const r3 = await p3 + + // calls the function with the latest provided args. + assert.deepStrictEqual((r3 as InlineCompletionItem[])[0].range?.end, new Position(2, 2)) + }) + }) }) }) }) diff --git a/packages/amazonq/test/unit/amazonq/apps/inline/inlineTracker.test.ts b/packages/amazonq/test/unit/amazonq/apps/inline/inlineTracker.test.ts new file mode 100644 index 00000000000..6b9490c72a5 --- /dev/null +++ b/packages/amazonq/test/unit/amazonq/apps/inline/inlineTracker.test.ts @@ -0,0 +1,299 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import { LineSelection, LineTracker, AuthUtil } from 'aws-core-vscode/codewhisperer' +import sinon from 'sinon' +import { Disposable, TextEditor, Position, Range, Selection } from 'vscode' +import { toTextEditor } from 'aws-core-vscode/test' +import assert from 'assert' +import { waitUntil } from 'aws-core-vscode/shared' + +describe('LineTracker class', function () { + let sut: LineTracker + let disposable: Disposable + let editor: TextEditor + let sandbox: sinon.SinonSandbox + let counts = { + editor: 0, + selection: 0, + content: 0, + } + + beforeEach(async function () { + sut = new LineTracker() + sandbox = sinon.createSandbox() + counts = { + editor: 0, + selection: 0, + content: 0, + } + disposable = sut.onDidChangeActiveLines((e) => { + if (e.reason === 'content') { + counts.content++ + } else if (e.reason === 'selection') { + counts.selection++ + } else if (e.reason === 'editor') { + counts.editor++ + } + }) + + sandbox.stub(AuthUtil.instance, 'isConnected').returns(true) + sandbox.stub(AuthUtil.instance, 'isConnectionExpired').returns(false) + }) + + afterEach(function () { + disposable.dispose() + sut.dispose() + sandbox.restore() + }) + + function assertEmptyCounts() { + assert.deepStrictEqual(counts, { + editor: 0, + selection: 0, + content: 0, + }) + } + + it('ready will emit onReady event', async function () { + let messageReceived = 0 + disposable = sut.onReady((_) => { + messageReceived++ + }) + + assert.strictEqual(sut.isReady, false) + sut.ready() + + await waitUntil( + async () => { + if (messageReceived !== 0) { + return + } + }, + { interval: 1000 } + ) + + assert.strictEqual(sut.isReady, true) + assert.strictEqual(messageReceived, 1) + }) + + describe('includes', function () { + // util function to help set up LineTracker.selections + async function setEditorSelection(selections: LineSelection[]): Promise { + const editor = await toTextEditor('\n\n\n\n\n\n\n\n\n\n', 'foo.py', undefined, { + preview: false, + }) + + const vscodeSelections = selections.map((s) => { + return new Selection(new Position(s.anchor, 0), new Position(s.active, 0)) + }) + + await sut.onTextEditorSelectionChanged({ + textEditor: editor, + selections: vscodeSelections, + kind: undefined, + }) + + assert.deepStrictEqual(sut.selections, selections) + return editor + } + + it('exact match when array of selections are provided', async function () { + const selections = [ + { + anchor: 1, + active: 1, + }, + { + anchor: 3, + active: 3, + }, + ] + + editor = await setEditorSelection(selections) + assert.deepStrictEqual(sut.selections, selections) + + let actual = sut.includes([ + { active: 1, anchor: 1 }, + { active: 3, anchor: 3 }, + ]) + assert.strictEqual(actual, true) + + actual = sut.includes([ + { active: 2, anchor: 2 }, + { active: 4, anchor: 4 }, + ]) + assert.strictEqual(actual, false) + + // both active && anchor have to be the same + actual = sut.includes([ + { active: 1, anchor: 0 }, + { active: 3, anchor: 0 }, + ]) + assert.strictEqual(actual, false) + + // different length would simply return false + actual = sut.includes([ + { active: 1, anchor: 1 }, + { active: 3, anchor: 3 }, + { active: 5, anchor: 5 }, + ]) + assert.strictEqual(actual, false) + }) + + it('match active line if line number and activeOnly option are provided', async function () { + const selections = [ + { + anchor: 1, + active: 1, + }, + { + anchor: 3, + active: 3, + }, + ] + + editor = await setEditorSelection(selections) + assert.deepStrictEqual(sut.selections, selections) + + let actual = sut.includes(1, { activeOnly: true }) + assert.strictEqual(actual, true) + + actual = sut.includes(2, { activeOnly: true }) + assert.strictEqual(actual, false) + }) + + it('range match if line number and activeOnly is set to false', async function () { + const selections = [ + { + anchor: 0, + active: 2, + }, + { + anchor: 4, + active: 6, + }, + ] + + editor = await setEditorSelection(selections) + assert.deepStrictEqual(sut.selections, selections) + + for (const line of [0, 1, 2]) { + const actual = sut.includes(line, { activeOnly: false }) + assert.strictEqual(actual, true) + } + + for (const line of [4, 5, 6]) { + const actual = sut.includes(line, { activeOnly: false }) + assert.strictEqual(actual, true) + } + + let actual = sut.includes(3, { activeOnly: false }) + assert.strictEqual(actual, false) + + actual = sut.includes(7, { activeOnly: false }) + assert.strictEqual(actual, false) + }) + }) + + describe('onContentChanged', function () { + it('should fire lineChangedEvent and set current line selection', async function () { + editor = await toTextEditor('\n\n\n\n\n', 'foo.py', undefined, { preview: false }) + editor.selection = new Selection(new Position(5, 0), new Position(5, 0)) + assertEmptyCounts() + + sut.onContentChanged({ + document: editor.document, + contentChanges: [{ text: 'a', range: new Range(0, 0, 0, 0), rangeOffset: 0, rangeLength: 0 }], + reason: undefined, + }) + + assert.deepStrictEqual(counts, { ...counts, content: 1 }) + assert.deepStrictEqual(sut.selections, [ + { + anchor: 5, + active: 5, + }, + ]) + }) + }) + + describe('onTextEditorSelectionChanged', function () { + it('should fire lineChangedEvent if selection changes and set current line selection', async function () { + editor = await toTextEditor('\n\n\n\n\n', 'foo.py', undefined, { preview: false }) + editor.selection = new Selection(new Position(3, 0), new Position(3, 0)) + assertEmptyCounts() + + await sut.onTextEditorSelectionChanged({ + textEditor: editor, + selections: [new Selection(new Position(3, 0), new Position(3, 0))], + kind: undefined, + }) + + assert.deepStrictEqual(counts, { ...counts, selection: 1 }) + assert.deepStrictEqual(sut.selections, [ + { + anchor: 3, + active: 3, + }, + ]) + + // if selection is included in the existing selections, won't emit an event + await sut.onTextEditorSelectionChanged({ + textEditor: editor, + selections: [new Selection(new Position(3, 0), new Position(3, 0))], + kind: undefined, + }) + + assert.deepStrictEqual(counts, { ...counts, selection: 1 }) + assert.deepStrictEqual(sut.selections, [ + { + anchor: 3, + active: 3, + }, + ]) + }) + + it('should not fire lineChangedEvent if uri scheme is debug || output', async function () { + // if the editor is not a text editor, won't emit an event and selection will be set to undefined + async function assertLineChanged(schema: string) { + const anotherEditor = await toTextEditor('', 'bar.log', undefined, { preview: false }) + const uri = anotherEditor.document.uri + sandbox.stub(uri, 'scheme').get(() => schema) + + await sut.onTextEditorSelectionChanged({ + textEditor: anotherEditor, + selections: [new Selection(new Position(3, 0), new Position(3, 0))], + kind: undefined, + }) + + assert.deepStrictEqual(counts, { ...counts }) + } + + await assertLineChanged('debug') + await assertLineChanged('output') + }) + }) + + describe('onActiveTextEditorChanged', function () { + it('shoudl fire lineChangedEvent', async function () { + const selections: Selection[] = [new Selection(0, 0, 1, 1)] + + editor = { selections: selections } as any + + assertEmptyCounts() + + await sut.onActiveTextEditorChanged(editor) + + assert.deepStrictEqual(counts, { ...counts, editor: 1 }) + assert.deepStrictEqual(sut.selections, [ + { + anchor: 0, + active: 1, + }, + ]) + }) + }) +}) diff --git a/packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts b/packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts index b3628e22c35..9bd195aa4a2 100644 --- a/packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts +++ b/packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts @@ -9,12 +9,19 @@ import { Position, CancellationToken, InlineCompletionItem } from 'vscode' import assert from 'assert' import { RecommendationService } from '../../../../../src/app/inline/recommendationService' import { SessionManager } from '../../../../../src/app/inline/sessionManager' -import { createMockDocument } from 'aws-core-vscode/test' +import { createMockDocument, createTestAuthUtil } from 'aws-core-vscode/test' +import { LineTracker } from '../../../../../src/app/inline/stateTracker/lineTracker' +import { InlineGeneratingMessage } from '../../../../../src/app/inline/inlineGeneratingMessage' describe('RecommendationService', () => { let languageClient: LanguageClient let sendRequestStub: sinon.SinonStub let sandbox: sinon.SinonSandbox + let sessionManager: SessionManager + let lineTracker: LineTracker + let activeStateController: InlineGeneratingMessage + let service: RecommendationService + const mockDocument = createMockDocument() const mockPosition = { line: 0, character: 0 } as Position const mockContext = { triggerKind: 1, selectedCompletionInfo: undefined } @@ -27,10 +34,8 @@ describe('RecommendationService', () => { insertText: 'ItemTwo', } as InlineCompletionItem const mockPartialResultToken = 'some-random-token' - const sessionManager = new SessionManager() - const service = new RecommendationService(sessionManager) - beforeEach(() => { + beforeEach(async () => { sandbox = sinon.createSandbox() sendRequestStub = sandbox.stub() @@ -38,6 +43,13 @@ describe('RecommendationService', () => { languageClient = { sendRequest: sendRequestStub, } as unknown as LanguageClient + + await createTestAuthUtil() + + sessionManager = new SessionManager() + lineTracker = new LineTracker() + activeStateController = new InlineGeneratingMessage(lineTracker) + service = new RecommendationService(sessionManager, activeStateController) }) afterEach(() => { @@ -107,13 +119,6 @@ describe('RecommendationService', () => { ...expectedRequestArgs, partialResultToken: mockPartialResultToken, }) - - // Verify session management - const items = sessionManager.getActiveRecommendation() - assert.deepStrictEqual(items, [mockInlineCompletionItemOne, { insertText: '1' } as InlineCompletionItem]) - sessionManager.incrementActiveIndex() - const items2 = sessionManager.getActiveRecommendation() - assert.deepStrictEqual(items2, [mockInlineCompletionItemTwo, { insertText: '1' } as InlineCompletionItem]) }) }) }) diff --git a/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts b/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts new file mode 100644 index 00000000000..5d9972019f4 --- /dev/null +++ b/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts @@ -0,0 +1,120 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import assert from 'assert' +import * as sinon from 'sinon' +import { assertTelemetry, createTestAuthUtil } from 'aws-core-vscode/test' +import { AuthUtil, awsIdSignIn, getStartUrl } from 'aws-core-vscode/codewhisperer' +import { backendAmazonQ } from 'aws-core-vscode/login' + +describe('Amazon Q Login', async function () { + const region = 'fakeRegion' + const startUrl = 'fakeUrl' + + let sandbox: sinon.SinonSandbox + let backend: backendAmazonQ.AmazonQLoginWebview + + beforeEach(async function () { + await createTestAuthUtil() + sandbox = sinon.createSandbox() + backend = new backendAmazonQ.AmazonQLoginWebview() + }) + + afterEach(function () { + sandbox.restore() + }) + + it('signs into builder ID and emits telemetry', async function () { + await backend.startBuilderIdSetup() + + assert.ok(AuthUtil.instance.isConnected()) + assert.ok(AuthUtil.instance.isBuilderIdConnection()) + + assertTelemetry('auth_addConnection', { + result: 'Succeeded', + credentialSourceId: 'awsId', + authEnabledFeatures: 'codewhisperer', + isReAuth: false, + ssoRegistrationExpiresAt: undefined, + ssoRegistrationClientId: undefined, + }) + }) + + it('signs into IdC and emits telemetry', async function () { + await backend.startEnterpriseSetup(startUrl, region) + + assert.ok(AuthUtil.instance.isConnected()) + assert.ok(AuthUtil.instance.isIdcConnection()) + assert.ok(AuthUtil.instance.isSsoSession()) + assert.deepStrictEqual(AuthUtil.instance.connection?.startUrl, startUrl) + assert.deepStrictEqual(AuthUtil.instance.connection?.region, region) + + assertTelemetry('auth_addConnection', { + result: 'Succeeded', + credentialSourceId: 'iamIdentityCenter', + authEnabledFeatures: 'codewhisperer', + credentialStartUrl: startUrl, + awsRegion: region, + isReAuth: false, + ssoRegistrationExpiresAt: undefined, + ssoRegistrationClientId: undefined, + }) + }) + + it('reauths builder ID and emits telemetry', async function () { + await awsIdSignIn() + + await backend.reauthenticateConnection() + + assert.ok(AuthUtil.instance.isConnected()) + + assertTelemetry('auth_addConnection', { + result: 'Succeeded', + credentialSourceId: 'awsId', + authEnabledFeatures: 'codewhisperer', + isReAuth: true, + ssoRegistrationExpiresAt: undefined, + ssoRegistrationClientId: undefined, + }) + }) + + it('reauths IdC and emits telemetry', async function () { + await getStartUrl.connectToEnterpriseSso(startUrl, region) + + await backend.reauthenticateConnection() + + assert.ok(AuthUtil.instance.isConnected()) + + assertTelemetry('auth_addConnection', { + result: 'Succeeded', + credentialSourceId: 'iamIdentityCenter', + authEnabledFeatures: 'codewhisperer', + credentialStartUrl: startUrl, + awsRegion: region, + isReAuth: true, + ssoRegistrationExpiresAt: undefined, + ssoRegistrationClientId: undefined, + }) + }) + + it('signs out of reauth and emits telemetry', async function () { + await getStartUrl.connectToEnterpriseSso(startUrl, region) + + await backend.signout() + + assert.ok(!AuthUtil.instance.isConnected()) + + assertTelemetry('auth_addConnection', { + result: 'Cancelled', + credentialSourceId: 'iamIdentityCenter', + authEnabledFeatures: 'codewhisperer', + credentialStartUrl: startUrl, + awsRegion: region, + isReAuth: true, + ssoRegistrationExpiresAt: undefined, + ssoRegistrationClientId: undefined, + }) + }) +}) diff --git a/packages/amazonq/test/unit/amazonq/lsp/auth.test.ts b/packages/amazonq/test/unit/amazonq/lsp/auth.test.ts deleted file mode 100644 index d55fef85f39..00000000000 --- a/packages/amazonq/test/unit/amazonq/lsp/auth.test.ts +++ /dev/null @@ -1,33 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -import assert from 'assert' -import { AmazonQLspAuth } from '../../../../src/lsp/auth' -import { LanguageClient } from 'vscode-languageclient' - -describe('AmazonQLspAuth', function () { - describe('updateBearerToken', function () { - it('makes request to LSP when token changes', async function () { - // Note: this token will be encrypted - let lastSentToken = {} - const auth = new AmazonQLspAuth({ - sendRequest: (_method: string, param: any) => { - lastSentToken = param - }, - info: (_message: string, _data: any) => {}, - } as LanguageClient) - - await auth.updateBearerToken('firstToken') - assert.notDeepStrictEqual(lastSentToken, {}) - const encryptedFirstToken = lastSentToken - - await auth.updateBearerToken('secondToken') - assert.notDeepStrictEqual(lastSentToken, encryptedFirstToken) - const encryptedSecondToken = lastSentToken - - await auth.updateBearerToken('secondToken') - assert.deepStrictEqual(lastSentToken, encryptedSecondToken) - }) - }) -}) diff --git a/packages/amazonq/test/unit/amazonq/lsp/chat/messages.test.ts b/packages/amazonq/test/unit/amazonq/lsp/chat/messages.test.ts index b2f5958f52b..caf74bda6e0 100644 --- a/packages/amazonq/test/unit/amazonq/lsp/chat/messages.test.ts +++ b/packages/amazonq/test/unit/amazonq/lsp/chat/messages.test.ts @@ -8,7 +8,7 @@ import { LanguageClient } from 'vscode-languageclient' import { AuthUtil } from 'aws-core-vscode/codewhisperer' import { registerMessageListeners } from '../../../../../src/lsp/chat/messages' import { AmazonQChatViewProvider } from '../../../../../src/lsp/chat/webviewProvider' -import { secondaryAuth, authConnection, AuthFollowUpType } from 'aws-core-vscode/amazonq' +import { AuthFollowUpType } from 'aws-core-vscode/amazonq' import { messages } from 'aws-core-vscode/shared' describe('registerMessageListeners', () => { @@ -50,7 +50,7 @@ describe('registerMessageListeners', () => { describe('AUTH_FOLLOW_UP_CLICKED', () => { let mockAuthUtil: AuthUtil - let deleteConnectionStub: sinon.SinonStub + let logoutStub: sinon.SinonStub let reauthenticateStub: sinon.SinonStub const authFollowUpClickedCommand = 'authFollowUpClicked' @@ -76,14 +76,12 @@ describe('registerMessageListeners', () => { } beforeEach(() => { - deleteConnectionStub = sandbox.stub().resolves() reauthenticateStub = sandbox.stub().resolves() + logoutStub = sandbox.stub().resolves() mockAuthUtil = { reauthenticate: reauthenticateStub, - secondaryAuth: { - deleteConnection: deleteConnectionStub, - } as unknown as secondaryAuth.SecondaryAuth, + logout: logoutStub, } as unknown as AuthUtil sandbox.replaceGetter(AuthUtil, 'instance', () => mockAuthUtil) @@ -98,7 +96,7 @@ describe('registerMessageListeners', () => { }) sinon.assert.calledOnce(reauthenticateStub) - sinon.assert.notCalled(deleteConnectionStub) + sinon.assert.notCalled(logoutStub) }) it('handles full authentication request', async () => { @@ -110,7 +108,7 @@ describe('registerMessageListeners', () => { }) sinon.assert.notCalled(reauthenticateStub) - sinon.assert.calledOnce(deleteConnectionStub) + sinon.assert.calledOnce(logoutStub) }) it('logs error if re-authentication fails', async () => { @@ -124,7 +122,7 @@ describe('registerMessageListeners', () => { it('logs error if full authentication fails', async () => { await testFailure({ authType: 'full-auth', - stubToReject: deleteConnectionStub, + stubToReject: logoutStub, errorMessage: 'Failed to authenticate', }) }) diff --git a/packages/amazonq/test/unit/amazonq/lsp/config.test.ts b/packages/amazonq/test/unit/amazonq/lsp/config.test.ts index 0327395fe1a..69b15d6e311 100644 --- a/packages/amazonq/test/unit/amazonq/lsp/config.test.ts +++ b/packages/amazonq/test/unit/amazonq/lsp/config.test.ts @@ -7,97 +7,73 @@ import assert from 'assert' import { DevSettings } from 'aws-core-vscode/shared' import sinon from 'sinon' import { defaultAmazonQLspConfig, ExtendedAmazonQLSPConfig, getAmazonQLspConfig } from '../../../../src/lsp/config' -import { defaultAmazonQWorkspaceLspConfig, getAmazonQWorkspaceLspConfig, LspConfig } from 'aws-core-vscode/amazonq' -for (const [name, config, defaultConfig, setEnv, resetEnv] of [ - [ - 'getAmazonQLspConfig', - getAmazonQLspConfig, - defaultAmazonQLspConfig, - (envConfig: ExtendedAmazonQLSPConfig) => { - process.env.__AMAZONQLSP_MANIFEST_URL = envConfig.manifestUrl - process.env.__AMAZONQLSP_SUPPORTED_VERSIONS = envConfig.supportedVersions - process.env.__AMAZONQLSP_ID = envConfig.id - process.env.__AMAZONQLSP_PATH = envConfig.path - process.env.__AMAZONQLSP_UI = envConfig.ui - }, - () => { - delete process.env.__AMAZONQLSP_MANIFEST_URL - delete process.env.__AMAZONQLSP_SUPPORTED_VERSIONS - delete process.env.__AMAZONQLSP_ID - delete process.env.__AMAZONQLSP_PATH - delete process.env.__AMAZONQLSP_UI - }, - ], - [ - 'getAmazonQWorkspaceLspConfig', - getAmazonQWorkspaceLspConfig, - defaultAmazonQWorkspaceLspConfig, - (envConfig: LspConfig) => { - process.env.__AMAZONQWORKSPACELSP_MANIFEST_URL = envConfig.manifestUrl - process.env.__AMAZONQWORKSPACELSP_SUPPORTED_VERSIONS = envConfig.supportedVersions - process.env.__AMAZONQWORKSPACELSP_ID = envConfig.id - process.env.__AMAZONQWORKSPACELSP_PATH = envConfig.path - }, - () => { - delete process.env.__AMAZONQWORKSPACELSP_MANIFEST_URL - delete process.env.__AMAZONQWORKSPACELSP_SUPPORTED_VERSIONS - delete process.env.__AMAZONQWORKSPACELSP_ID - delete process.env.__AMAZONQWORKSPACELSP_PATH - }, - ], -] as const) { - describe(name, () => { - let sandbox: sinon.SinonSandbox - let serviceConfigStub: sinon.SinonStub - const settingConfig: LspConfig = { - manifestUrl: 'https://custom.url/manifest.json', - supportedVersions: '4.0.0', - id: 'AmazonQSetting', - suppressPromptPrefix: config().suppressPromptPrefix, - path: '/custom/path', - ...(name === 'getAmazonQLspConfig' && { ui: '/chat/client/location' }), - } +describe('getAmazonQLspConfig', () => { + let sandbox: sinon.SinonSandbox + let serviceConfigStub: sinon.SinonStub + const settingConfig: ExtendedAmazonQLSPConfig = { + manifestUrl: 'https://custom.url/manifest.json', + supportedVersions: '4.0.0', + id: 'AmazonQSetting', + suppressPromptPrefix: getAmazonQLspConfig().suppressPromptPrefix, + path: '/custom/path', + ui: '/chat/client/location', + } - beforeEach(() => { - sandbox = sinon.createSandbox() + beforeEach(() => { + sandbox = sinon.createSandbox() - serviceConfigStub = sandbox.stub() - sandbox.stub(DevSettings, 'instance').get(() => ({ - getServiceConfig: serviceConfigStub, - })) - }) + serviceConfigStub = sandbox.stub() + sandbox.stub(DevSettings, 'instance').get(() => ({ + getServiceConfig: serviceConfigStub, + })) + }) - afterEach(() => { - sandbox.restore() - resetEnv() - }) + afterEach(() => { + sandbox.restore() + resetEnv() + }) - it('uses default config', () => { - serviceConfigStub.returns({}) - assert.deepStrictEqual(config(), defaultConfig) - }) + it('uses default config', () => { + serviceConfigStub.returns({}) + assert.deepStrictEqual(getAmazonQLspConfig(), defaultAmazonQLspConfig) + }) - it('overrides path', () => { - const path = '/custom/path/to/lsp' - serviceConfigStub.returns({ path }) + it('overrides path', () => { + const path = '/custom/path/to/lsp' + serviceConfigStub.returns({ path }) - assert.deepStrictEqual(config(), { - ...defaultConfig, - path, - }) + assert.deepStrictEqual(getAmazonQLspConfig(), { + ...defaultAmazonQLspConfig, + path, }) + }) - it('overrides default settings', () => { - serviceConfigStub.returns(settingConfig) + it('overrides default settings', () => { + serviceConfigStub.returns(settingConfig) - assert.deepStrictEqual(config(), settingConfig) - }) + assert.deepStrictEqual(getAmazonQLspConfig(), settingConfig) + }) - it('environment variable takes precedence over settings', () => { - setEnv(settingConfig) - serviceConfigStub.returns({}) - assert.deepStrictEqual(config(), settingConfig) - }) + it('environment variable takes precedence over settings', () => { + setEnv(settingConfig) + serviceConfigStub.returns({}) + assert.deepStrictEqual(getAmazonQLspConfig(), settingConfig) }) -} + + function setEnv(envConfig: ExtendedAmazonQLSPConfig) { + process.env.__AMAZONQLSP_MANIFEST_URL = envConfig.manifestUrl + process.env.__AMAZONQLSP_SUPPORTED_VERSIONS = envConfig.supportedVersions + process.env.__AMAZONQLSP_ID = envConfig.id + process.env.__AMAZONQLSP_PATH = envConfig.path + process.env.__AMAZONQLSP_UI = envConfig.ui + } + + function resetEnv() { + delete process.env.__AMAZONQLSP_MANIFEST_URL + delete process.env.__AMAZONQLSP_SUPPORTED_VERSIONS + delete process.env.__AMAZONQLSP_ID + delete process.env.__AMAZONQLSP_PATH + delete process.env.__AMAZONQLSP_UI + } +}) diff --git a/packages/amazonq/test/unit/amazonq/lsp/encryption.test.ts b/packages/amazonq/test/unit/amazonq/lsp/encryption.test.ts new file mode 100644 index 00000000000..78cf5351a98 --- /dev/null +++ b/packages/amazonq/test/unit/amazonq/lsp/encryption.test.ts @@ -0,0 +1,27 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as assert from 'assert' +import { decryptResponse, encryptRequest } from '../../../../src/lsp/encryption' +import { encryptionKey } from '../../../../src/lsp/client' + +describe('LSP encryption', function () { + it('encrypt and decrypt invert eachother with same key', async function () { + const key = encryptionKey + const request = { + id: 0, + name: 'my Request', + isRealRequest: false, + metadata: { + tags: ['tag1', 'tag2'], + }, + } + const encryptedPayload = await encryptRequest(request, key) + const message = (encryptedPayload as { message: string }).message + const decrypted = await decryptResponse(message, key) + + assert.deepStrictEqual(decrypted, request) + }) +}) diff --git a/packages/amazonq/test/unit/amazonq/lsp/lspClient.test.ts b/packages/amazonq/test/unit/amazonq/lsp/lspClient.test.ts deleted file mode 100644 index 369cda5402d..00000000000 --- a/packages/amazonq/test/unit/amazonq/lsp/lspClient.test.ts +++ /dev/null @@ -1,72 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -import * as sinon from 'sinon' -import assert from 'assert' -import { globals, getNodeExecutableName } from 'aws-core-vscode/shared' -import { LspClient, lspClient as lspClientModule } from 'aws-core-vscode/amazonq' - -describe('Amazon Q LSP client', function () { - let lspClient: LspClient - let encryptFunc: sinon.SinonSpy - - beforeEach(async function () { - sinon.stub(globals, 'isWeb').returns(false) - lspClient = new LspClient() - encryptFunc = sinon.spy(lspClient, 'encrypt') - }) - - it('encrypts payload of query ', async () => { - await lspClient.queryVectorIndex('mock_input') - assert.ok(encryptFunc.calledOnce) - assert.ok(encryptFunc.calledWith(JSON.stringify({ query: 'mock_input' }))) - const value = await encryptFunc.returnValues[0] - // verifies JWT encryption header - assert.ok(value.startsWith(`eyJhbGciOiJkaXIiLCJlbmMiOiJBMjU2R0NNIn0`)) - }) - - it('encrypts payload of index files ', async () => { - await lspClient.buildIndex(['fileA'], 'path', 'all') - assert.ok(encryptFunc.calledOnce) - assert.ok( - encryptFunc.calledWith( - JSON.stringify({ - filePaths: ['fileA'], - projectRoot: 'path', - config: 'all', - language: '', - }) - ) - ) - const value = await encryptFunc.returnValues[0] - // verifies JWT encryption header - assert.ok(value.startsWith(`eyJhbGciOiJkaXIiLCJlbmMiOiJBMjU2R0NNIn0`)) - }) - - it('encrypt removes readable information', async () => { - const sample = 'hello' - const encryptedSample = await lspClient.encrypt(sample) - assert.ok(!encryptedSample.includes('hello')) - }) - - it('validates node executable + lsp bundle', async () => { - await assert.rejects(async () => { - await lspClientModule.activate(globals.context, { - // Mimic the `LspResolution` type. - node: 'node.bogus.exe', - lsp: 'fake/lsp.js', - }) - }, /.*failed to run basic .*node.*exitcode.*node\.bogus\.exe.*/) - await assert.rejects(async () => { - await lspClientModule.activate(globals.context, { - node: getNodeExecutableName(), - lsp: 'fake/lsp.js', - }) - }, /.*failed to run .*exitcode.*node.*lsp\.js/) - }) - - afterEach(() => { - sinon.restore() - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/commands/invokeRecommendation.test.ts b/packages/amazonq/test/unit/codewhisperer/commands/invokeRecommendation.test.ts deleted file mode 100644 index 68cebe37bb1..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/commands/invokeRecommendation.test.ts +++ /dev/null @@ -1,43 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as sinon from 'sinon' -import { resetCodeWhispererGlobalVariables, createMockTextEditor } from 'aws-core-vscode/test' -import { - ConfigurationEntry, - invokeRecommendation, - InlineCompletionService, - isInlineCompletionEnabled, - DefaultCodeWhispererClient, -} from 'aws-core-vscode/codewhisperer' - -describe('invokeRecommendation', function () { - describe('invokeRecommendation', function () { - let getRecommendationStub: sinon.SinonStub - let mockClient: DefaultCodeWhispererClient - - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - getRecommendationStub = sinon.stub(InlineCompletionService.instance, 'getPaginatedRecommendation') - }) - - afterEach(function () { - sinon.restore() - }) - - it('Should call getPaginatedRecommendation with OnDemand as trigger type when inline completion is enabled', async function () { - const mockEditor = createMockTextEditor() - const config: ConfigurationEntry = { - isShowMethodsEnabled: true, - isManualTriggerEnabled: true, - isAutomatedTriggerEnabled: true, - isSuggestionsWithCodeReferencesEnabled: true, - } - await invokeRecommendation(mockEditor, mockClient, config) - assert.strictEqual(getRecommendationStub.called, isInlineCompletionEnabled()) - }) - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/commands/onAcceptance.test.ts b/packages/amazonq/test/unit/codewhisperer/commands/onAcceptance.test.ts deleted file mode 100644 index 0471aaa3601..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/commands/onAcceptance.test.ts +++ /dev/null @@ -1,64 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as vscode from 'vscode' -import * as sinon from 'sinon' -import { onAcceptance, AcceptedSuggestionEntry, session, CodeWhispererTracker } from 'aws-core-vscode/codewhisperer' -import { resetCodeWhispererGlobalVariables, createMockTextEditor } from 'aws-core-vscode/test' - -describe('onAcceptance', function () { - describe('onAcceptance', function () { - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - session.reset() - }) - - afterEach(function () { - sinon.restore() - session.reset() - }) - - it('Should enqueue an event object to tracker', async function () { - const mockEditor = createMockTextEditor() - const trackerSpy = sinon.spy(CodeWhispererTracker.prototype, 'enqueue') - const fakeReferences = [ - { - message: '', - licenseName: 'MIT', - repository: 'http://github.com/fake', - recommendationContentSpan: { - start: 0, - end: 10, - }, - }, - ] - await onAcceptance({ - editor: mockEditor, - range: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 26)), - effectiveRange: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 26)), - acceptIndex: 0, - recommendation: "print('Hello World!')", - requestId: '', - sessionId: '', - triggerType: 'OnDemand', - completionType: 'Line', - language: 'python', - references: fakeReferences, - }) - const actualArg = trackerSpy.getCall(0).args[0] as AcceptedSuggestionEntry - assert.ok(trackerSpy.calledOnce) - assert.strictEqual(actualArg.originalString, 'def two_sum(nums, target):') - assert.strictEqual(actualArg.requestId, '') - assert.strictEqual(actualArg.sessionId, '') - assert.strictEqual(actualArg.triggerType, 'OnDemand') - assert.strictEqual(actualArg.completionType, 'Line') - assert.strictEqual(actualArg.language, 'python') - assert.deepStrictEqual(actualArg.startPosition, new vscode.Position(1, 0)) - assert.deepStrictEqual(actualArg.endPosition, new vscode.Position(1, 26)) - assert.strictEqual(actualArg.index, 0) - }) - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/commands/onInlineAcceptance.test.ts b/packages/amazonq/test/unit/codewhisperer/commands/onInlineAcceptance.test.ts deleted file mode 100644 index ed3bc99fa34..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/commands/onInlineAcceptance.test.ts +++ /dev/null @@ -1,43 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as vscode from 'vscode' -import * as sinon from 'sinon' -import { resetCodeWhispererGlobalVariables, createMockTextEditor } from 'aws-core-vscode/test' -import { onInlineAcceptance, RecommendationHandler, session } from 'aws-core-vscode/codewhisperer' - -describe('onInlineAcceptance', function () { - describe('onInlineAcceptance', function () { - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - session.reset() - }) - - afterEach(function () { - sinon.restore() - session.reset() - }) - - it('Should dispose inline completion provider', async function () { - const mockEditor = createMockTextEditor() - const spy = sinon.spy(RecommendationHandler.instance, 'disposeInlineCompletion') - await onInlineAcceptance({ - editor: mockEditor, - range: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 21)), - effectiveRange: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 21)), - acceptIndex: 0, - recommendation: "print('Hello World!')", - requestId: '', - sessionId: '', - triggerType: 'OnDemand', - completionType: 'Line', - language: 'python', - references: undefined, - }) - assert.ok(spy.calledWith()) - }) - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts b/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts index 11441b9bf6f..a77e47e33ab 100644 --- a/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts @@ -7,41 +7,38 @@ import * as sinon from 'sinon' import assert, { fail } from 'assert' import { AuthUtil, RegionProfile, RegionProfileManager, defaultServiceConfig } from 'aws-core-vscode/codewhisperer' import { globals } from 'aws-core-vscode/shared' -import { createTestAuth } from 'aws-core-vscode/test' -import { SsoConnection } from 'aws-core-vscode/auth' +import { constants } from 'aws-core-vscode/auth' +import { createTestAuthUtil } from 'aws-core-vscode/test' +import { randomUUID } from 'crypto' const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start' +const region = 'us-east-1' -describe('RegionProfileManager', function () { - let sut: RegionProfileManager - let auth: ReturnType - let authUtil: AuthUtil +describe('RegionProfileManager', async function () { + let regionProfileManager: RegionProfileManager const profileFoo: RegionProfile = { name: 'foo', - region: 'us-east-1', + region, arn: 'foo arn', description: 'foo description', } async function setupConnection(type: 'builderId' | 'idc') { if (type === 'builderId') { - await authUtil.connectToAwsBuilderId() - const conn = authUtil.conn - assert.strictEqual(conn?.type, 'sso') - assert.strictEqual(conn.label, 'AWS Builder ID') + await AuthUtil.instance.login(constants.builderIdStartUrl, region) + assert.ok(AuthUtil.instance.isSsoSession()) + assert.ok(AuthUtil.instance.isBuilderIdConnection()) } else if (type === 'idc') { - await authUtil.connectToEnterpriseSso(enterpriseSsoStartUrl, 'us-east-1') - const conn = authUtil.conn - assert.strictEqual(conn?.type, 'sso') - assert.strictEqual(conn.label, 'IAM Identity Center (enterprise)') + await AuthUtil.instance.login(enterpriseSsoStartUrl, region) + assert.ok(AuthUtil.instance.isSsoSession()) + assert.ok(AuthUtil.instance.isIdcConnection()) } } - beforeEach(function () { - auth = createTestAuth(globals.globalState) - authUtil = new AuthUtil(auth) - sut = new RegionProfileManager(() => authUtil.conn) + beforeEach(async function () { + await createTestAuthUtil() + regionProfileManager = new RegionProfileManager(AuthUtil.instance) }) afterEach(function () { @@ -65,12 +62,12 @@ describe('RegionProfileManager', function () { const mockClient = { listAvailableProfiles: listProfilesStub, } - const createClientStub = sinon.stub(sut, '_createQClient').resolves(mockClient) + const createClientStub = sinon.stub(regionProfileManager, '_createQClient').resolves(mockClient) - const r = await sut.listRegionProfile() + const profileList = await regionProfileManager.listRegionProfile() - assert.strictEqual(r.length, 2) - assert.deepStrictEqual(r, [ + assert.strictEqual(profileList.length, 2) + assert.deepStrictEqual(profileList, [ { name: 'foo', arn: 'arn', @@ -93,41 +90,42 @@ describe('RegionProfileManager', function () { describe('switch and get profile', function () { it('should switch if connection is IdC', async function () { await setupConnection('idc') - await sut.switchRegionProfile(profileFoo, 'user') - assert.deepStrictEqual(sut.activeRegionProfile, profileFoo) + await regionProfileManager.switchRegionProfile(profileFoo, 'user') + assert.deepStrictEqual(regionProfileManager.activeRegionProfile, profileFoo) }) it('should do nothing and return undefined if connection is builder id', async function () { await setupConnection('builderId') - await sut.switchRegionProfile(profileFoo, 'user') - assert.deepStrictEqual(sut.activeRegionProfile, undefined) + await regionProfileManager.switchRegionProfile(profileFoo, 'user') + assert.deepStrictEqual(regionProfileManager.activeRegionProfile, undefined) }) }) describe(`client config`, function () { it(`no valid credential should throw`, async function () { - assert.ok(authUtil.conn === undefined) + await AuthUtil.instance.logout() + + assert.ok(!AuthUtil.instance.isConnected()) assert.throws(() => { - sut.clientConfig + regionProfileManager.clientConfig }, /trying to get client configuration without credential/) }) it(`builder id should always use default profile IAD`, async function () { await setupConnection('builderId') - await sut.switchRegionProfile(profileFoo, 'user') - assert.deepStrictEqual(sut.activeRegionProfile, undefined) - const conn = authUtil.conn - if (!conn) { + await regionProfileManager.switchRegionProfile(profileFoo, 'user') + assert.deepStrictEqual(regionProfileManager.activeRegionProfile, undefined) + if (!AuthUtil.instance.isConnected()) { fail('connection should not be undefined') } - assert.deepStrictEqual(sut.clientConfig, defaultServiceConfig) + assert.deepStrictEqual(regionProfileManager.clientConfig, defaultServiceConfig) }) it(`idc should return correct endpoint corresponding to profile region`, async function () { await setupConnection('idc') - await sut.switchRegionProfile( + await regionProfileManager.switchRegionProfile( { name: 'foo', region: 'eu-central-1', @@ -136,8 +134,8 @@ describe('RegionProfileManager', function () { }, 'user' ) - assert.ok(sut.activeRegionProfile) - assert.deepStrictEqual(sut.clientConfig, { + assert.ok(regionProfileManager.activeRegionProfile) + assert.deepStrictEqual(regionProfileManager.clientConfig, { region: 'eu-central-1', endpoint: 'https://q.eu-central-1.amazonaws.com/', }) @@ -145,7 +143,7 @@ describe('RegionProfileManager', function () { it(`idc should throw if corresponding endpoint is not defined`, async function () { await setupConnection('idc') - await sut.switchRegionProfile( + await regionProfileManager.switchRegionProfile( { name: 'foo', region: 'unknown region', @@ -156,22 +154,21 @@ describe('RegionProfileManager', function () { ) assert.throws(() => { - sut.clientConfig + regionProfileManager.clientConfig }, /Q client configuration error, endpoint not found for region*/) }) }) - describe('persistence', function () { + describe('persistSelectedRegionProfile', function () { it('persistSelectedRegionProfile', async function () { await setupConnection('idc') - await sut.switchRegionProfile(profileFoo, 'user') - assert.deepStrictEqual(sut.activeRegionProfile, profileFoo) - const conn = authUtil.conn - if (!conn) { + await regionProfileManager.switchRegionProfile(profileFoo, 'user') + assert.deepStrictEqual(regionProfileManager.activeRegionProfile, profileFoo) + if (!AuthUtil.instance.isConnected()) { fail('connection should not be undefined') } - await sut.persistSelectRegionProfile() + await regionProfileManager.persistSelectRegionProfile() const state = globals.globalState.tryGet<{ [label: string]: RegionProfile }>( 'aws.amazonq.regionProfiles', @@ -179,25 +176,68 @@ describe('RegionProfileManager', function () { {} ) - assert.strictEqual(state[conn.id], profileFoo) + assert.strictEqual(state[AuthUtil.instance.profileName], profileFoo) }) + }) - it(`restoreRegionProfile`, async function () { - sinon.stub(sut, 'listRegionProfile').resolves([profileFoo]) + describe('restoreRegionProfile', function () { + beforeEach(async function () { await setupConnection('idc') - const conn = authUtil.conn - if (!conn) { - fail('connection should not be undefined') - } + }) + it('restores region profile if profile name matches', async function () { + const state = {} as any + state[AuthUtil.instance.profileName] = profileFoo + + await globals.globalState.update('aws.amazonq.regionProfiles', state) + + await regionProfileManager.restoreRegionProfile() + + assert.strictEqual(regionProfileManager.activeRegionProfile, profileFoo) + }) + + it('returns early when no profiles exist', async function () { + const state = {} as any + state[AuthUtil.instance.profileName] = undefined + + await globals.globalState.update('aws.amazonq.regionProfiles', state) + + await regionProfileManager.restoreRegionProfile() + assert.strictEqual(regionProfileManager.activeRegionProfile, undefined) + }) + + it('returns early when no profile name matches, and multiple profiles exist', async function () { + const state = {} as any + state[AuthUtil.instance.profileName] = undefined + state[randomUUID()] = profileFoo + + await globals.globalState.update('aws.amazonq.regionProfiles', state) + await regionProfileManager.restoreRegionProfile() + assert.strictEqual(regionProfileManager.activeRegionProfile, undefined) + }) + + it('uses single profile when no profile name matches', async function () { const state = {} as any - state[conn.id] = profileFoo + state[randomUUID()] = profileFoo await globals.globalState.update('aws.amazonq.regionProfiles', state) - await sut.restoreRegionProfile(conn) + await regionProfileManager.restoreRegionProfile() + + assert.strictEqual(regionProfileManager.activeRegionProfile, profileFoo) + }) + + it('handles cross-validation failure', async function () { + const state = { + [AuthUtil.instance.profileName]: profileFoo, + } + sinon.stub(regionProfileManager, 'loadPersistedRegionProfiles').returns(state) + sinon.stub(regionProfileManager, 'getProfiles').resolves([]) // No matching profile + const invalidateStub = sinon.stub(regionProfileManager, 'invalidateProfile') + + await regionProfileManager.restoreRegionProfile() - assert.strictEqual(sut.activeRegionProfile, profileFoo) + assert.ok(invalidateStub.calledWith(profileFoo.arn)) }) }) @@ -205,25 +245,24 @@ describe('RegionProfileManager', function () { it('should reset activeProfile and global state', async function () { // setup await setupConnection('idc') - await sut.switchRegionProfile(profileFoo, 'user') - assert.deepStrictEqual(sut.activeRegionProfile, profileFoo) - const conn = authUtil.conn - if (!conn) { + await regionProfileManager.switchRegionProfile(profileFoo, 'user') + assert.deepStrictEqual(regionProfileManager.activeRegionProfile, profileFoo) + if (!AuthUtil.instance.isConnected()) { fail('connection should not be undefined') } - await sut.persistSelectRegionProfile() + await regionProfileManager.persistSelectRegionProfile() const state = globals.globalState.tryGet<{ [label: string]: RegionProfile }>( 'aws.amazonq.regionProfiles', Object, {} ) - assert.strictEqual(state[conn.id], profileFoo) + assert.strictEqual(state[AuthUtil.instance.profileName], profileFoo) // subject to test - await sut.invalidateProfile(profileFoo.arn) + await regionProfileManager.invalidateProfile(profileFoo.arn) // assertion - assert.strictEqual(sut.activeRegionProfile, undefined) + assert.strictEqual(regionProfileManager.activeRegionProfile, undefined) const actualGlobalState = globals.globalState.tryGet<{ [label: string]: RegionProfile }>( 'aws.amazonq.regionProfiles', Object, @@ -237,7 +276,7 @@ describe('RegionProfileManager', function () { it(`should configure the endpoint and region from a profile`, async function () { await setupConnection('idc') - const iadClient = await sut.createQClient({ + const iadClient = await regionProfileManager.createQClient({ name: 'foo', region: 'us-east-1', arn: 'arn', @@ -247,7 +286,7 @@ describe('RegionProfileManager', function () { assert.deepStrictEqual(iadClient.config.region, 'us-east-1') assert.deepStrictEqual(iadClient.endpoint.href, 'https://q.us-east-1.amazonaws.com/') - const fraClient = await sut.createQClient({ + const fraClient = await regionProfileManager.createQClient({ name: 'bar', region: 'eu-central-1', arn: 'arn', @@ -263,7 +302,7 @@ describe('RegionProfileManager', function () { await assert.rejects( async () => { - await sut.createQClient({ + await regionProfileManager.createQClient({ name: 'foo', region: 'ap-east-1', arn: 'arn', @@ -275,7 +314,7 @@ describe('RegionProfileManager', function () { await assert.rejects( async () => { - await sut.createQClient({ + await regionProfileManager.createQClient({ name: 'foo', region: 'unknown-somewhere', arn: 'arn', @@ -288,11 +327,10 @@ describe('RegionProfileManager', function () { it(`should configure the endpoint and region correspondingly`, async function () { await setupConnection('idc') - await sut.switchRegionProfile(profileFoo, 'user') - assert.deepStrictEqual(sut.activeRegionProfile, profileFoo) - const conn = authUtil.conn as SsoConnection + await regionProfileManager.switchRegionProfile(profileFoo, 'user') + assert.deepStrictEqual(regionProfileManager.activeRegionProfile, profileFoo) - const client = await sut._createQClient('eu-central-1', 'https://amazon.com/', conn) + const client = await regionProfileManager._createQClient('eu-central-1', 'https://amazon.com/') assert.deepStrictEqual(client.config.region, 'eu-central-1') assert.deepStrictEqual(client.endpoint.href, 'https://amazon.com/') diff --git a/packages/amazonq/test/unit/codewhisperer/service/codewhisperer.test.ts b/packages/amazonq/test/unit/codewhisperer/service/codewhisperer.test.ts index f30d92de496..e222fb9bda4 100644 --- a/packages/amazonq/test/unit/codewhisperer/service/codewhisperer.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/service/codewhisperer.test.ts @@ -136,7 +136,7 @@ describe('codewhisperer', async function () { }), } as Request) - const authUtilStub = sinon.stub(AuthUtil.instance, 'isValidEnterpriseSsoInUse').returns(isSso) + const authUtilStub = sinon.stub(AuthUtil.instance, 'isIdcConnection').returns(isSso) await globals.telemetry.setTelemetryEnabled(isTelemetryEnabled) await codeWhispererClient.sendTelemetryEvent({ telemetryEvent: payload }) const expectedOptOutPreference = isTelemetryEnabled ? 'OPTIN' : 'OPTOUT' diff --git a/packages/amazonq/test/unit/codewhisperer/service/completionProvider.test.ts b/packages/amazonq/test/unit/codewhisperer/service/completionProvider.test.ts deleted file mode 100644 index 956999d64ad..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/service/completionProvider.test.ts +++ /dev/null @@ -1,117 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as vscode from 'vscode' - -import { - getCompletionItems, - getCompletionItem, - getLabel, - Recommendation, - RecommendationHandler, - session, -} from 'aws-core-vscode/codewhisperer' -import { createMockDocument, resetCodeWhispererGlobalVariables } from 'aws-core-vscode/test' - -describe('completionProviderService', function () { - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - }) - - describe('getLabel', function () { - it('should return correct label given recommendation longer than Constants.LABEL_LENGTH', function () { - const mockLongRecommendation = ` - const metaDataFile = path.join(__dirname, 'nls.metadata.json'); - const locale = getUserDefinedLocale(argvConfig);` - const expected = '\n const m..' - assert.strictEqual(getLabel(mockLongRecommendation), expected) - }) - - it('should return correct label given short recommendation', function () { - const mockShortRecommendation = 'function onReady()' - const expected = 'function onReady()..' - assert.strictEqual(getLabel(mockShortRecommendation), expected) - }) - }) - - describe('getCompletionItem', function () { - it('should return targetCompletionItem given input', function () { - session.startPos = new vscode.Position(0, 0) - RecommendationHandler.instance.requestId = 'mock_requestId_getCompletionItem' - session.sessionId = 'mock_sessionId_getCompletionItem' - const mockPosition = new vscode.Position(0, 1) - const mockRecommendationDetail: Recommendation = { - content: "\n\t\tconsole.log('Hello world!');\n\t}", - } - const mockRecommendationIndex = 1 - const mockDocument = createMockDocument('', 'test.ts', 'typescript') - const expected: vscode.CompletionItem = { - label: "\n\t\tconsole.log('Hell..", - kind: 1, - detail: 'CodeWhisperer', - documentation: new vscode.MarkdownString().appendCodeblock( - "\n\t\tconsole.log('Hello world!');\n\t}", - 'typescript' - ), - sortText: '0000000002', - preselect: true, - insertText: new vscode.SnippetString("\n\t\tconsole.log('Hello world!');\n\t}"), - keepWhitespace: true, - command: { - command: 'aws.amazonq.accept', - title: 'On acceptance', - arguments: [ - new vscode.Range(0, 0, 0, 0), - 1, - "\n\t\tconsole.log('Hello world!');\n\t}", - 'mock_requestId_getCompletionItem', - 'mock_sessionId_getCompletionItem', - 'OnDemand', - 'Line', - 'typescript', - undefined, - ], - }, - } - const actual = getCompletionItem( - mockDocument, - mockPosition, - mockRecommendationDetail, - mockRecommendationIndex - ) - assert.deepStrictEqual(actual.command, expected.command) - assert.strictEqual(actual.sortText, expected.sortText) - assert.strictEqual(actual.label, expected.label) - assert.strictEqual(actual.kind, expected.kind) - assert.strictEqual(actual.preselect, expected.preselect) - assert.strictEqual(actual.keepWhitespace, expected.keepWhitespace) - assert.strictEqual(JSON.stringify(actual.documentation), JSON.stringify(expected.documentation)) - assert.strictEqual(JSON.stringify(actual.insertText), JSON.stringify(expected.insertText)) - }) - }) - - describe('getCompletionItems', function () { - it('should return completion items for each non-empty recommendation', async function () { - session.recommendations = [ - { content: "\n\t\tconsole.log('Hello world!');\n\t}" }, - { content: '\nvar a = 10' }, - ] - const mockPosition = new vscode.Position(0, 0) - const mockDocument = createMockDocument('', 'test.ts', 'typescript') - const actual = getCompletionItems(mockDocument, mockPosition) - assert.strictEqual(actual.length, 2) - }) - - it('should return empty completion items when recommendation is empty', async function () { - session.recommendations = [] - const mockPosition = new vscode.Position(14, 83) - const mockDocument = createMockDocument() - const actual = getCompletionItems(mockDocument, mockPosition) - const expected: vscode.CompletionItem[] = [] - assert.deepStrictEqual(actual, expected) - }) - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts b/packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts deleted file mode 100644 index 18fd7d2f21b..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts +++ /dev/null @@ -1,255 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import assert from 'assert' -import * as sinon from 'sinon' -import { - CodeWhispererStatusBar, - InlineCompletionService, - ReferenceInlineProvider, - RecommendationHandler, - CodeSuggestionsState, - ConfigurationEntry, - CWInlineCompletionItemProvider, - session, - AuthUtil, - listCodeWhispererCommandsId, - DefaultCodeWhispererClient, -} from 'aws-core-vscode/codewhisperer' -import { createMockTextEditor, resetCodeWhispererGlobalVariables, createMockDocument } from 'aws-core-vscode/test' - -describe('inlineCompletionService', function () { - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - }) - - describe('getPaginatedRecommendation', function () { - const config: ConfigurationEntry = { - isShowMethodsEnabled: true, - isManualTriggerEnabled: true, - isAutomatedTriggerEnabled: true, - isSuggestionsWithCodeReferencesEnabled: true, - } - - let mockClient: DefaultCodeWhispererClient - - beforeEach(async function () { - mockClient = new DefaultCodeWhispererClient() - await resetCodeWhispererGlobalVariables() - }) - - afterEach(function () { - sinon.restore() - }) - - it('should call checkAndResetCancellationTokens before showing inline and next token to be null', async function () { - const mockEditor = createMockTextEditor() - sinon.stub(RecommendationHandler.instance, 'getRecommendations').resolves({ - result: 'Succeeded', - errorMessage: undefined, - recommendationCount: 1, - }) - const checkAndResetCancellationTokensStub = sinon.stub( - RecommendationHandler.instance, - 'checkAndResetCancellationTokens' - ) - session.recommendations = [{ content: "\n\t\tconsole.log('Hello world!');\n\t}" }, { content: '' }] - await InlineCompletionService.instance.getPaginatedRecommendation( - mockClient, - mockEditor, - 'OnDemand', - config - ) - assert.ok(checkAndResetCancellationTokensStub.called) - assert.strictEqual(RecommendationHandler.instance.hasNextToken(), false) - }) - }) - - describe('clearInlineCompletionStates', function () { - it('should remove inline reference and recommendations', async function () { - const fakeReferences = [ - { - message: '', - licenseName: 'MIT', - repository: 'http://github.com/fake', - recommendationContentSpan: { - start: 0, - end: 10, - }, - }, - ] - ReferenceInlineProvider.instance.setInlineReference(1, 'test', fakeReferences) - session.recommendations = [{ content: "\n\t\tconsole.log('Hello world!');\n\t}" }, { content: '' }] - session.language = 'python' - - assert.ok(session.recommendations.length > 0) - await RecommendationHandler.instance.clearInlineCompletionStates() - assert.strictEqual(ReferenceInlineProvider.instance.refs.length, 0) - assert.strictEqual(session.recommendations.length, 0) - }) - }) - - describe('truncateOverlapWithRightContext', function () { - const fileName = 'test.py' - const language = 'python' - const rightContext = 'return target\n' - const doc = `import math\ndef two_sum(nums, target):\n` - const provider = new CWInlineCompletionItemProvider(0, 0, [], '', new vscode.Position(0, 0), '') - - it('removes overlap with right context from suggestion', async function () { - const mockSuggestion = 'return target\n' - const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language) - const cursorPosition = new vscode.Position(2, 0) - const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion, cursorPosition) - assert.strictEqual(result, '') - }) - - it('only removes the overlap part from suggestion', async function () { - const mockSuggestion = 'print(nums)\nreturn target\n' - const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language) - const cursorPosition = new vscode.Position(2, 0) - const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion, cursorPosition) - assert.strictEqual(result, 'print(nums)\n') - }) - - it('only removes the last overlap pattern from suggestion', async function () { - const mockSuggestion = 'return target\nprint(nums)\nreturn target\n' - const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language) - const cursorPosition = new vscode.Position(2, 0) - const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion, cursorPosition) - assert.strictEqual(result, 'return target\nprint(nums)\n') - }) - - it('returns empty string if the remaining suggestion only contains white space', async function () { - const mockSuggestion = 'return target\n ' - const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language) - const cursorPosition = new vscode.Position(2, 0) - const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion, cursorPosition) - assert.strictEqual(result, '') - }) - - it('returns the original suggestion if no match found', async function () { - const mockSuggestion = 'import numpy\n' - const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language) - const cursorPosition = new vscode.Position(2, 0) - const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion, cursorPosition) - assert.strictEqual(result, 'import numpy\n') - }) - - it('ignores the space at the end of recommendation', async function () { - const mockSuggestion = 'return target\n\n\n\n\n' - const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language) - const cursorPosition = new vscode.Position(2, 0) - const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion, cursorPosition) - assert.strictEqual(result, '') - }) - }) -}) - -describe('CWInlineCompletionProvider', function () { - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - }) - - describe('provideInlineCompletionItems', function () { - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - }) - - afterEach(function () { - sinon.restore() - }) - - it('should return undefined if position is before RecommendationHandler start pos', async function () { - const position = new vscode.Position(0, 0) - const document = createMockDocument() - const fakeContext = { triggerKind: 0, selectedCompletionInfo: undefined } - const token = new vscode.CancellationTokenSource().token - const provider = new CWInlineCompletionItemProvider(0, 0, [], '', new vscode.Position(1, 1), '') - const result = await provider.provideInlineCompletionItems(document, position, fakeContext, token) - - assert.ok(result === undefined) - }) - }) -}) - -describe('codewhisperer status bar', function () { - let sandbox: sinon.SinonSandbox - let statusBar: TestStatusBar - let service: InlineCompletionService - - class TestStatusBar extends CodeWhispererStatusBar { - constructor() { - super() - } - - getStatusBar() { - return this.statusBar - } - } - - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - sandbox = sinon.createSandbox() - statusBar = new TestStatusBar() - service = new InlineCompletionService(statusBar) - }) - - afterEach(function () { - sandbox.restore() - }) - - it('shows correct status bar when auth is not connected', async function () { - sandbox.stub(AuthUtil.instance, 'isConnectionValid').returns(false) - sandbox.stub(AuthUtil.instance, 'isConnectionExpired').returns(false) - - await service.refreshStatusBar() - - const actualStatusBar = statusBar.getStatusBar() - assert.strictEqual(actualStatusBar.text, '$(chrome-close) Amazon Q') - assert.strictEqual(actualStatusBar.command, listCodeWhispererCommandsId) - assert.deepStrictEqual(actualStatusBar.backgroundColor, new vscode.ThemeColor('statusBarItem.errorBackground')) - }) - - it('shows correct status bar when auth is connected', async function () { - sandbox.stub(AuthUtil.instance, 'isConnectionValid').returns(true) - sandbox.stub(CodeSuggestionsState.instance, 'isSuggestionsEnabled').returns(true) - - await service.refreshStatusBar() - - const actualStatusBar = statusBar.getStatusBar() - assert.strictEqual(actualStatusBar.text, '$(debug-start) Amazon Q') - assert.strictEqual(actualStatusBar.command, listCodeWhispererCommandsId) - assert.deepStrictEqual(actualStatusBar.backgroundColor, undefined) - }) - - it('shows correct status bar when auth is connected but paused', async function () { - sandbox.stub(AuthUtil.instance, 'isConnectionValid').returns(true) - sandbox.stub(CodeSuggestionsState.instance, 'isSuggestionsEnabled').returns(false) - - await service.refreshStatusBar() - - const actualStatusBar = statusBar.getStatusBar() - assert.strictEqual(actualStatusBar.text, '$(debug-pause) Amazon Q') - assert.strictEqual(actualStatusBar.command, listCodeWhispererCommandsId) - assert.deepStrictEqual(actualStatusBar.backgroundColor, undefined) - }) - - it('shows correct status bar when auth is expired', async function () { - sandbox.stub(AuthUtil.instance, 'isConnectionValid').returns(false) - sandbox.stub(AuthUtil.instance, 'isConnectionExpired').returns(true) - - await service.refreshStatusBar() - - const actualStatusBar = statusBar.getStatusBar() - assert.strictEqual(actualStatusBar.text, '$(debug-disconnect) Amazon Q') - assert.strictEqual(actualStatusBar.command, listCodeWhispererCommandsId) - assert.deepStrictEqual( - actualStatusBar.backgroundColor, - new vscode.ThemeColor('statusBarItem.warningBackground') - ) - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/service/keyStrokeHandler.test.ts b/packages/amazonq/test/unit/codewhisperer/service/keyStrokeHandler.test.ts deleted file mode 100644 index 4b6a5291f22..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/service/keyStrokeHandler.test.ts +++ /dev/null @@ -1,237 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as vscode from 'vscode' -import * as sinon from 'sinon' -import * as codewhispererSdkClient from 'aws-core-vscode/codewhisperer' -import { - createMockTextEditor, - createTextDocumentChangeEvent, - resetCodeWhispererGlobalVariables, -} from 'aws-core-vscode/test' -import * as EditorContext from 'aws-core-vscode/codewhisperer' -import { - ConfigurationEntry, - DocumentChangedSource, - KeyStrokeHandler, - DefaultDocumentChangedType, - RecommendationService, - ClassifierTrigger, - isInlineCompletionEnabled, - RecommendationHandler, - InlineCompletionService, -} from 'aws-core-vscode/codewhisperer' - -describe('keyStrokeHandler', function () { - const config: ConfigurationEntry = { - isShowMethodsEnabled: true, - isManualTriggerEnabled: true, - isAutomatedTriggerEnabled: true, - isSuggestionsWithCodeReferencesEnabled: true, - } - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - }) - describe('processKeyStroke', async function () { - let invokeSpy: sinon.SinonStub - let startTimerSpy: sinon.SinonStub - let mockClient: codewhispererSdkClient.DefaultCodeWhispererClient - beforeEach(async function () { - invokeSpy = sinon.stub(KeyStrokeHandler.instance, 'invokeAutomatedTrigger') - startTimerSpy = sinon.stub(KeyStrokeHandler.instance, 'startIdleTimeTriggerTimer') - sinon.spy(RecommendationHandler.instance, 'getRecommendations') - mockClient = new codewhispererSdkClient.DefaultCodeWhispererClient() - await resetCodeWhispererGlobalVariables() - sinon.stub(mockClient, 'listRecommendations') - sinon.stub(mockClient, 'generateRecommendations') - }) - afterEach(function () { - sinon.restore() - }) - - it('Whatever the input is, should skip when automatic trigger is turned off, should not call invokeAutomatedTrigger', async function () { - const mockEditor = createMockTextEditor() - const mockEvent: vscode.TextDocumentChangeEvent = createTextDocumentChangeEvent( - mockEditor.document, - new vscode.Range(new vscode.Position(0, 0), new vscode.Position(0, 1)), - ' ' - ) - const cfg: ConfigurationEntry = { - isShowMethodsEnabled: true, - isManualTriggerEnabled: true, - isAutomatedTriggerEnabled: false, - isSuggestionsWithCodeReferencesEnabled: true, - } - const keyStrokeHandler = new KeyStrokeHandler() - await keyStrokeHandler.processKeyStroke(mockEvent, mockEditor, mockClient, cfg) - assert.ok(!invokeSpy.called) - assert.ok(!startTimerSpy.called) - }) - - it('Should not call invokeAutomatedTrigger when changed text across multiple lines', async function () { - await testShouldInvoke('\nprint(n', false) - }) - - it('Should not call invokeAutomatedTrigger when doing delete or undo (empty changed text)', async function () { - await testShouldInvoke('', false) - }) - - it('Should call invokeAutomatedTrigger with Enter when inputing \n', async function () { - await testShouldInvoke('\n', true) - }) - - it('Should call invokeAutomatedTrigger with Enter when inputing \r\n', async function () { - await testShouldInvoke('\r\n', true) - }) - - it('Should call invokeAutomatedTrigger with SpecialCharacter when inputing {', async function () { - await testShouldInvoke('{', true) - }) - - it('Should not call invokeAutomatedTrigger for non-special characters for classifier language if classifier says no', async function () { - sinon.stub(ClassifierTrigger.instance, 'shouldTriggerFromClassifier').returns(false) - await testShouldInvoke('a', false) - }) - - it('Should call invokeAutomatedTrigger for non-special characters for classifier language if classifier says yes', async function () { - sinon.stub(ClassifierTrigger.instance, 'shouldTriggerFromClassifier').returns(true) - await testShouldInvoke('a', true) - }) - - it('Should skip invoking if there is immediate right context on the same line and not a single }', async function () { - const casesForSuppressTokenFilling = [ - { - rightContext: 'add', - shouldInvoke: false, - }, - { - rightContext: '}', - shouldInvoke: true, - }, - { - rightContext: '} ', - shouldInvoke: true, - }, - { - rightContext: ')', - shouldInvoke: true, - }, - { - rightContext: ') ', - shouldInvoke: true, - }, - { - rightContext: ' add', - shouldInvoke: true, - }, - { - rightContext: ' ', - shouldInvoke: true, - }, - { - rightContext: '\naddTwo', - shouldInvoke: true, - }, - ] - - for (const o of casesForSuppressTokenFilling) { - await testShouldInvoke('{', o.shouldInvoke, o.rightContext) - } - }) - - async function testShouldInvoke(input: string, shouldTrigger: boolean, rightContext: string = '') { - const mockEditor = createMockTextEditor(rightContext, 'test.js', 'javascript', 0, 0) - const mockEvent: vscode.TextDocumentChangeEvent = createTextDocumentChangeEvent( - mockEditor.document, - new vscode.Range(new vscode.Position(0, 0), new vscode.Position(0, 1)), - input - ) - await KeyStrokeHandler.instance.processKeyStroke(mockEvent, mockEditor, mockClient, config) - assert.strictEqual( - invokeSpy.called, - shouldTrigger, - `invokeAutomatedTrigger ${shouldTrigger ? 'NOT' : 'WAS'} called for rightContext: "${rightContext}"` - ) - } - }) - - describe('invokeAutomatedTrigger', function () { - let mockClient: codewhispererSdkClient.DefaultCodeWhispererClient - beforeEach(async function () { - sinon.restore() - mockClient = new codewhispererSdkClient.DefaultCodeWhispererClient() - await resetCodeWhispererGlobalVariables() - sinon.stub(mockClient, 'listRecommendations') - sinon.stub(mockClient, 'generateRecommendations') - }) - afterEach(function () { - sinon.restore() - }) - - it('should call getPaginatedRecommendation when inline completion is enabled', async function () { - const mockEditor = createMockTextEditor() - const keyStrokeHandler = new KeyStrokeHandler() - const mockEvent: vscode.TextDocumentChangeEvent = createTextDocumentChangeEvent( - mockEditor.document, - new vscode.Range(new vscode.Position(0, 0), new vscode.Position(0, 1)), - ' ' - ) - const getRecommendationsStub = sinon.stub(InlineCompletionService.instance, 'getPaginatedRecommendation') - await keyStrokeHandler.invokeAutomatedTrigger('Enter', mockEditor, mockClient, config, mockEvent) - assert.strictEqual(getRecommendationsStub.called, isInlineCompletionEnabled()) - }) - }) - - describe('shouldTriggerIdleTime', function () { - it('should return false when inline is enabled and inline completion is in progress ', function () { - const keyStrokeHandler = new KeyStrokeHandler() - sinon.stub(RecommendationService.instance, 'isRunning').get(() => true) - const result = keyStrokeHandler.shouldTriggerIdleTime() - assert.strictEqual(result, !isInlineCompletionEnabled()) - }) - }) - - describe('test checkChangeSource', function () { - const tabStr = ' '.repeat(EditorContext.getTabSize()) - - const cases: [string, DocumentChangedSource][] = [ - ['\n ', DocumentChangedSource.EnterKey], - ['\n', DocumentChangedSource.EnterKey], - ['(', DocumentChangedSource.SpecialCharsKey], - ['()', DocumentChangedSource.SpecialCharsKey], - ['{}', DocumentChangedSource.SpecialCharsKey], - ['(a, b):', DocumentChangedSource.Unknown], - [':', DocumentChangedSource.SpecialCharsKey], - ['a', DocumentChangedSource.RegularKey], - [tabStr, DocumentChangedSource.TabKey], - [' ', DocumentChangedSource.Reformatting], - ['def add(a,b):\n return a + b\n', DocumentChangedSource.Unknown], - ['function suggestedByIntelliSense():', DocumentChangedSource.Unknown], - ] - - for (const tuple of cases) { - const input = tuple[0] - const expected = tuple[1] - it(`test input ${input} should return ${expected}`, function () { - const actual = new DefaultDocumentChangedType( - createFakeDocumentChangeEvent(tuple[0]) - ).checkChangeSource() - assert.strictEqual(actual, expected) - }) - } - - function createFakeDocumentChangeEvent(str: string): ReadonlyArray { - return [ - { - range: new vscode.Range(new vscode.Position(0, 0), new vscode.Position(0, 5)), - rangeOffset: 0, - rangeLength: 0, - text: str, - }, - ] - } - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/service/recommendationHandler.test.ts b/packages/amazonq/test/unit/codewhisperer/service/recommendationHandler.test.ts deleted file mode 100644 index 86dfc5e514c..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/service/recommendationHandler.test.ts +++ /dev/null @@ -1,271 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as vscode from 'vscode' -import * as sinon from 'sinon' -import { - ReferenceInlineProvider, - session, - AuthUtil, - DefaultCodeWhispererClient, - RecommendationsList, - ConfigurationEntry, - RecommendationHandler, - CodeWhispererCodeCoverageTracker, - supplementalContextUtil, -} from 'aws-core-vscode/codewhisperer' -import { - assertTelemetryCurried, - stub, - createMockTextEditor, - resetCodeWhispererGlobalVariables, -} from 'aws-core-vscode/test' -// import * as supplementalContextUtil from 'aws-core-vscode/codewhisperer' - -describe('recommendationHandler', function () { - const config: ConfigurationEntry = { - isShowMethodsEnabled: true, - isManualTriggerEnabled: true, - isAutomatedTriggerEnabled: true, - isSuggestionsWithCodeReferencesEnabled: true, - } - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - }) - - describe('getRecommendations', async function () { - const mockClient = stub(DefaultCodeWhispererClient) - const mockEditor = createMockTextEditor() - const testStartUrl = 'testStartUrl' - - beforeEach(async function () { - sinon.restore() - await resetCodeWhispererGlobalVariables() - mockClient.listRecommendations.resolves({}) - mockClient.generateRecommendations.resolves({}) - RecommendationHandler.instance.clearRecommendations() - sinon.stub(AuthUtil.instance, 'startUrl').value(testStartUrl) - }) - - afterEach(function () { - sinon.restore() - }) - - it('should assign correct recommendations given input', async function () { - assert.strictEqual(CodeWhispererCodeCoverageTracker.instances.size, 0) - assert.strictEqual( - CodeWhispererCodeCoverageTracker.getTracker(mockEditor.document.languageId)?.serviceInvocationCount, - 0 - ) - - const mockServerResult = { - recommendations: [{ content: "print('Hello World!')" }, { content: '' }], - $response: { - requestId: 'test_request', - httpResponse: { - headers: { - 'x-amzn-sessionid': 'test_request', - }, - }, - }, - } - const handler = new RecommendationHandler() - sinon.stub(handler, 'getServerResponse').resolves(mockServerResult) - await handler.getRecommendations(mockClient, mockEditor, 'AutoTrigger', config, 'Enter', false) - const actual = session.recommendations - const expected: RecommendationsList = [{ content: "print('Hello World!')" }, { content: '' }] - assert.deepStrictEqual(actual, expected) - assert.strictEqual( - CodeWhispererCodeCoverageTracker.getTracker(mockEditor.document.languageId)?.serviceInvocationCount, - 1 - ) - }) - - it('should assign request id correctly', async function () { - const mockServerResult = { - recommendations: [{ content: "print('Hello World!')" }, { content: '' }], - $response: { - requestId: 'test_request', - httpResponse: { - headers: { - 'x-amzn-sessionid': 'test_request', - }, - }, - }, - } - const handler = new RecommendationHandler() - sinon.stub(handler, 'getServerResponse').resolves(mockServerResult) - sinon.stub(handler, 'isCancellationRequested').returns(false) - await handler.getRecommendations(mockClient, mockEditor, 'AutoTrigger', config, 'Enter', false) - assert.strictEqual(handler.requestId, 'test_request') - assert.strictEqual(session.sessionId, 'test_request') - assert.strictEqual(session.triggerType, 'AutoTrigger') - }) - - it('should call telemetry function that records a CodeWhisperer service invocation', async function () { - const mockServerResult = { - recommendations: [{ content: "print('Hello World!')" }, { content: '' }], - $response: { - requestId: 'test_request', - httpResponse: { - headers: { - 'x-amzn-sessionid': 'test_request', - }, - }, - }, - } - const handler = new RecommendationHandler() - sinon.stub(handler, 'getServerResponse').resolves(mockServerResult) - sinon.stub(supplementalContextUtil, 'fetchSupplementalContext').resolves({ - isUtg: false, - isProcessTimeout: false, - supplementalContextItems: [], - contentsLength: 100, - latency: 0, - strategy: 'empty', - }) - sinon.stub(performance, 'now').returns(0.0) - session.startPos = new vscode.Position(1, 0) - session.startCursorOffset = 2 - await handler.getRecommendations(mockClient, mockEditor, 'AutoTrigger', config, 'Enter') - const assertTelemetry = assertTelemetryCurried('codewhisperer_serviceInvocation') - assertTelemetry({ - codewhispererRequestId: 'test_request', - codewhispererSessionId: 'test_request', - codewhispererLastSuggestionIndex: 1, - codewhispererTriggerType: 'AutoTrigger', - codewhispererAutomatedTriggerType: 'Enter', - codewhispererImportRecommendationEnabled: true, - result: 'Succeeded', - codewhispererLineNumber: 1, - codewhispererCursorOffset: 38, - codewhispererLanguage: 'python', - credentialStartUrl: testStartUrl, - codewhispererSupplementalContextIsUtg: false, - codewhispererSupplementalContextTimeout: false, - codewhispererSupplementalContextLatency: 0, - codewhispererSupplementalContextLength: 100, - }) - }) - }) - - describe('isValidResponse', function () { - afterEach(function () { - sinon.restore() - }) - it('should return true if any response is not empty', function () { - const handler = new RecommendationHandler() - session.recommendations = [ - { - content: - '\n // Use the console to output debug info…n of the command with the "command" variable', - }, - { content: '' }, - ] - assert.ok(handler.isValidResponse()) - }) - - it('should return false if response is empty', function () { - const handler = new RecommendationHandler() - session.recommendations = [] - assert.ok(!handler.isValidResponse()) - }) - - it('should return false if all response has no string length', function () { - const handler = new RecommendationHandler() - session.recommendations = [{ content: '' }, { content: '' }] - assert.ok(!handler.isValidResponse()) - }) - }) - - describe('setCompletionType/getCompletionType', function () { - beforeEach(function () { - sinon.restore() - }) - - it('should set the completion type to block given a multi-line suggestion', function () { - session.setCompletionType(0, { content: 'test\n\n \t\r\nanother test' }) - assert.strictEqual(session.getCompletionType(0), 'Block') - - session.setCompletionType(0, { content: 'test\ntest\n' }) - assert.strictEqual(session.getCompletionType(0), 'Block') - - session.setCompletionType(0, { content: '\n \t\r\ntest\ntest' }) - assert.strictEqual(session.getCompletionType(0), 'Block') - }) - - it('should set the completion type to line given a single-line suggestion', function () { - session.setCompletionType(0, { content: 'test' }) - assert.strictEqual(session.getCompletionType(0), 'Line') - - session.setCompletionType(0, { content: 'test\r\t ' }) - assert.strictEqual(session.getCompletionType(0), 'Line') - }) - - it('should set the completion type to line given a multi-line completion but only one-lien of non-blank sequence', function () { - session.setCompletionType(0, { content: 'test\n\t' }) - assert.strictEqual(session.getCompletionType(0), 'Line') - - session.setCompletionType(0, { content: 'test\n ' }) - assert.strictEqual(session.getCompletionType(0), 'Line') - - session.setCompletionType(0, { content: 'test\n\r' }) - assert.strictEqual(session.getCompletionType(0), 'Line') - - session.setCompletionType(0, { content: '\n\n\n\ntest' }) - assert.strictEqual(session.getCompletionType(0), 'Line') - }) - }) - - describe('on event change', async function () { - beforeEach(function () { - const fakeReferences = [ - { - message: '', - licenseName: 'MIT', - repository: 'http://github.com/fake', - recommendationContentSpan: { - start: 0, - end: 10, - }, - }, - ] - ReferenceInlineProvider.instance.setInlineReference(1, 'test', fakeReferences) - session.sessionId = '' - RecommendationHandler.instance.requestId = '' - }) - - it('should remove inline reference onEditorChange', async function () { - session.sessionId = 'aSessionId' - RecommendationHandler.instance.requestId = 'aRequestId' - await RecommendationHandler.instance.onEditorChange() - assert.strictEqual(ReferenceInlineProvider.instance.refs.length, 0) - }) - it('should remove inline reference onFocusChange', async function () { - session.sessionId = 'aSessionId' - RecommendationHandler.instance.requestId = 'aRequestId' - await RecommendationHandler.instance.onFocusChange() - assert.strictEqual(ReferenceInlineProvider.instance.refs.length, 0) - }) - it('should not remove inline reference on cursor change from typing', async function () { - await RecommendationHandler.instance.onCursorChange({ - textEditor: createMockTextEditor(), - selections: [], - kind: vscode.TextEditorSelectionChangeKind.Keyboard, - }) - assert.strictEqual(ReferenceInlineProvider.instance.refs.length, 1) - }) - - it('should remove inline reference on cursor change from mouse movement', async function () { - await RecommendationHandler.instance.onCursorChange({ - textEditor: vscode.window.activeTextEditor!, - selections: [], - kind: vscode.TextEditorSelectionChangeKind.Mouse, - }) - assert.strictEqual(ReferenceInlineProvider.instance.refs.length, 0) - }) - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/service/referenceLogViewProvider.test.ts b/packages/amazonq/test/unit/codewhisperer/service/referenceLogViewProvider.test.ts index 1c1b6322675..dcacf745a57 100644 --- a/packages/amazonq/test/unit/codewhisperer/service/referenceLogViewProvider.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/service/referenceLogViewProvider.test.ts @@ -5,7 +5,6 @@ import assert from 'assert' import { createMockTextEditor, resetCodeWhispererGlobalVariables } from 'aws-core-vscode/test' import { ReferenceLogViewProvider, LicenseUtil } from 'aws-core-vscode/codewhisperer' - describe('referenceLogViewProvider', function () { beforeEach(async function () { await resetCodeWhispererGlobalVariables() @@ -66,4 +65,39 @@ describe('referenceLogViewProvider', function () { assert.ok(!actual.includes(LicenseUtil.getLicenseHtml('MIT'))) }) }) + + it('accepts references from CW and language server', async function () { + const cwReference = { + licenseName: 'MIT', + repository: 'TEST_REPO', + url: 'cw.com', + recommendationContentSpan: { + start: 0, + end: 10, + }, + } + + const flareReference = { + referenceName: 'test reference', + referenceUrl: 'flare.com', + licenseName: 'apache', + position: { + startCharacter: 0, + endCharacter: 10, + }, + } + + const actual = ReferenceLogViewProvider.getReferenceLog( + '', + [cwReference, flareReference], + createMockTextEditor() + ) + + assert.ok(actual.includes('MIT')) + assert.ok(actual.includes('apache')) + assert.ok(actual.includes('TEST_REPO')) + assert.ok(actual.includes('test reference')) + assert.ok(actual.includes('flare.com')) + assert.ok(actual.includes('cw.com')) + }) }) diff --git a/packages/amazonq/test/unit/codewhisperer/service/telemetry.test.ts b/packages/amazonq/test/unit/codewhisperer/service/telemetry.test.ts index 0f1429f130b..1f11661f002 100644 --- a/packages/amazonq/test/unit/codewhisperer/service/telemetry.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/service/telemetry.test.ts @@ -19,9 +19,7 @@ import { DefaultCodeWhispererClient, ListRecommendationsResponse, Recommendation, - invokeRecommendation, ConfigurationEntry, - RecommendationHandler, session, vsCodeCursorUpdateDelay, AuthUtil, @@ -113,7 +111,6 @@ describe.skip('CodeWhisperer telemetry', async function () { }) async function resetStates() { - await RecommendationHandler.instance.clearInlineCompletionStates() await resetCodeWhispererGlobalVariables() } @@ -424,7 +421,6 @@ describe.skip('CodeWhisperer telemetry', async function () { assert.strictEqual(session.sessionId, 'session_id_1') assert.deepStrictEqual(session.requestIdList, ['request_id_1', 'request_id_1', 'request_id_1_2']) - await RecommendationHandler.instance.onEditorChange() assertSessionClean() await backspace(editor) // todo: without this, the following manual trigger will not be displayed in the test, investigate and fix it @@ -500,7 +496,6 @@ describe.skip('CodeWhisperer telemetry', async function () { await manualTrigger(editor, client, config) await assertTextEditorContains('') - await RecommendationHandler.instance.onFocusChange() assertTelemetry('codewhisperer_userTriggerDecision', [ session1UserTriggerEvent({ codewhispererSuggestionState: 'Reject' }), ]) @@ -513,7 +508,6 @@ async function manualTrigger( client: DefaultCodeWhispererClient, config: ConfigurationEntry ) { - await invokeRecommendation(editor, client, config) await waitUntilSuggestionSeen() } diff --git a/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererCodeCoverageTracker.test.ts b/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererCodeCoverageTracker.test.ts deleted file mode 100644 index ee001b3328d..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererCodeCoverageTracker.test.ts +++ /dev/null @@ -1,560 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as sinon from 'sinon' -import * as vscode from 'vscode' -import { - CodeWhispererCodeCoverageTracker, - vsCodeState, - TelemetryHelper, - AuthUtil, - getUnmodifiedAcceptedTokens, -} from 'aws-core-vscode/codewhisperer' -import { createMockDocument, createMockTextEditor, resetCodeWhispererGlobalVariables } from 'aws-core-vscode/test' -import { globals } from 'aws-core-vscode/shared' -import { assertTelemetryCurried } from 'aws-core-vscode/test' - -describe('codewhispererCodecoverageTracker', function () { - const language = 'python' - - describe('test getTracker', function () { - afterEach(async function () { - await resetCodeWhispererGlobalVariables() - CodeWhispererCodeCoverageTracker.instances.clear() - }) - - it('unsupported language', function () { - assert.strictEqual(CodeWhispererCodeCoverageTracker.getTracker('vb'), undefined) - assert.strictEqual(CodeWhispererCodeCoverageTracker.getTracker('ipynb'), undefined) - }) - - it('supported language', function () { - assert.notStrictEqual(CodeWhispererCodeCoverageTracker.getTracker('python'), undefined) - assert.notStrictEqual(CodeWhispererCodeCoverageTracker.getTracker('javascriptreact'), undefined) - assert.notStrictEqual(CodeWhispererCodeCoverageTracker.getTracker('java'), undefined) - assert.notStrictEqual(CodeWhispererCodeCoverageTracker.getTracker('javascript'), undefined) - assert.notStrictEqual(CodeWhispererCodeCoverageTracker.getTracker('cpp'), undefined) - assert.notStrictEqual(CodeWhispererCodeCoverageTracker.getTracker('ruby'), undefined) - assert.notStrictEqual(CodeWhispererCodeCoverageTracker.getTracker('go'), undefined) - }) - - it('supported language and should return singleton object per language', function () { - let instance1: CodeWhispererCodeCoverageTracker | undefined - let instance2: CodeWhispererCodeCoverageTracker | undefined - instance1 = CodeWhispererCodeCoverageTracker.getTracker('java') - instance2 = CodeWhispererCodeCoverageTracker.getTracker('java') - assert.notStrictEqual(instance1, undefined) - assert.strictEqual(Object.is(instance1, instance2), true) - - instance1 = CodeWhispererCodeCoverageTracker.getTracker('python') - instance2 = CodeWhispererCodeCoverageTracker.getTracker('python') - assert.notStrictEqual(instance1, undefined) - assert.strictEqual(Object.is(instance1, instance2), true) - - instance1 = CodeWhispererCodeCoverageTracker.getTracker('javascriptreact') - instance2 = CodeWhispererCodeCoverageTracker.getTracker('javascriptreact') - assert.notStrictEqual(instance1, undefined) - assert.strictEqual(Object.is(instance1, instance2), true) - }) - }) - - describe('test isActive', function () { - let tracker: CodeWhispererCodeCoverageTracker | undefined - - afterEach(async function () { - await resetCodeWhispererGlobalVariables() - CodeWhispererCodeCoverageTracker.instances.clear() - sinon.restore() - }) - - it('inactive case: telemetryEnable = true, isConnected = false', function () { - sinon.stub(TelemetryHelper.instance, 'isTelemetryEnabled').returns(true) - sinon.stub(AuthUtil.instance, 'isConnected').returns(false) - - tracker = CodeWhispererCodeCoverageTracker.getTracker('python') - if (!tracker) { - assert.fail() - } - - assert.strictEqual(tracker.isActive(), false) - }) - - it('inactive case: telemetryEnabled = false, isConnected = false', function () { - sinon.stub(TelemetryHelper.instance, 'isTelemetryEnabled').returns(false) - sinon.stub(AuthUtil.instance, 'isConnected').returns(false) - - tracker = CodeWhispererCodeCoverageTracker.getTracker('java') - if (!tracker) { - assert.fail() - } - - assert.strictEqual(tracker.isActive(), false) - }) - - it('active case: telemetryEnabled = true, isConnected = true', function () { - sinon.stub(TelemetryHelper.instance, 'isTelemetryEnabled').returns(true) - sinon.stub(AuthUtil.instance, 'isConnected').returns(true) - - tracker = CodeWhispererCodeCoverageTracker.getTracker('javascript') - if (!tracker) { - assert.fail() - } - assert.strictEqual(tracker.isActive(), true) - }) - }) - - describe('updateAcceptedTokensCount', function () { - let tracker: CodeWhispererCodeCoverageTracker | undefined - - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - tracker = CodeWhispererCodeCoverageTracker.getTracker(language) - if (tracker) { - sinon.stub(tracker, 'isActive').returns(true) - } - }) - - afterEach(function () { - sinon.restore() - CodeWhispererCodeCoverageTracker.instances.clear() - }) - - it('Should compute edit distance to update the accepted tokens', function () { - if (!tracker) { - assert.fail() - } - const editor = createMockTextEditor('def addTwoNumbers(a, b):\n') - - tracker.addAcceptedTokens(editor.document.fileName, { - range: new vscode.Range(0, 0, 0, 25), - text: `def addTwoNumbers(x, y):\n`, - accepted: 25, - }) - tracker.addTotalTokens(editor.document.fileName, 100) - tracker.updateAcceptedTokensCount(editor) - assert.strictEqual(tracker?.acceptedTokens[editor.document.fileName][0].accepted, 23) - }) - }) - - describe('getUnmodifiedAcceptedTokens', function () { - let tracker: CodeWhispererCodeCoverageTracker | undefined - - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - tracker = CodeWhispererCodeCoverageTracker.getTracker(language) - if (tracker) { - sinon.stub(tracker, 'isActive').returns(true) - } - }) - - afterEach(function () { - sinon.restore() - CodeWhispererCodeCoverageTracker.instances.clear() - }) - - it('Should return correct unmodified accepted tokens count', function () { - assert.strictEqual(getUnmodifiedAcceptedTokens('foo', 'fou'), 2) - assert.strictEqual(getUnmodifiedAcceptedTokens('foo', 'f11111oo'), 3) - assert.strictEqual(getUnmodifiedAcceptedTokens('foo', 'fo'), 2) - assert.strictEqual(getUnmodifiedAcceptedTokens('helloworld', 'HelloWorld'), 8) - assert.strictEqual(getUnmodifiedAcceptedTokens('helloworld', 'World'), 4) - assert.strictEqual(getUnmodifiedAcceptedTokens('CodeWhisperer', 'CODE'), 1) - assert.strictEqual(getUnmodifiedAcceptedTokens('CodeWhisperer', 'CodeWhispererGood'), 13) - }) - }) - - describe('countAcceptedTokens', function () { - let tracker: CodeWhispererCodeCoverageTracker | undefined - - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - tracker = CodeWhispererCodeCoverageTracker.getTracker(language) - if (tracker) { - sinon.stub(tracker, 'isActive').returns(true) - } - }) - - afterEach(function () { - sinon.restore() - CodeWhispererCodeCoverageTracker.instances.clear() - }) - - it('Should skip when tracker is not active', function () { - if (!tracker) { - assert.fail() - } - tracker.countAcceptedTokens(new vscode.Range(0, 0, 0, 1), 'a', 'test.py') - const spy = sinon.spy(CodeWhispererCodeCoverageTracker.prototype, 'addAcceptedTokens') - assert.ok(!spy.called) - }) - - it('Should increase AcceptedTokens', function () { - if (!tracker) { - assert.fail() - } - tracker.countAcceptedTokens(new vscode.Range(0, 0, 0, 1), 'a', 'test.py') - assert.deepStrictEqual(tracker.acceptedTokens['test.py'][0], { - range: new vscode.Range(0, 0, 0, 1), - text: 'a', - accepted: 1, - }) - }) - it('Should increase TotalTokens', function () { - if (!tracker) { - assert.fail() - } - tracker.countAcceptedTokens(new vscode.Range(0, 0, 0, 1), 'a', 'test.py') - tracker.countAcceptedTokens(new vscode.Range(0, 0, 0, 1), 'b', 'test.py') - assert.deepStrictEqual(tracker.totalTokens['test.py'], 2) - }) - }) - - describe('countTotalTokens', function () { - let tracker: CodeWhispererCodeCoverageTracker | undefined - - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - tracker = CodeWhispererCodeCoverageTracker.getTracker(language) - if (tracker) { - sinon.stub(tracker, 'isActive').returns(true) - } - }) - - afterEach(function () { - sinon.restore() - CodeWhispererCodeCoverageTracker.instances.clear() - }) - - it('Should skip when content change size is more than 50', function () { - if (!tracker) { - assert.fail() - } - tracker.countTotalTokens({ - reason: undefined, - document: createMockDocument(), - contentChanges: [ - { - range: new vscode.Range(0, 0, 0, 600), - rangeOffset: 0, - rangeLength: 600, - text: 'def twoSum(nums, target):\nfor '.repeat(20), - }, - ], - }) - assert.strictEqual(Object.keys(tracker.totalTokens).length, 0) - }) - - it('Should not skip when content change size is less than 50', function () { - if (!tracker) { - assert.fail() - } - tracker.countTotalTokens({ - reason: undefined, - document: createMockDocument(), - contentChanges: [ - { - range: new vscode.Range(0, 0, 0, 49), - rangeOffset: 0, - rangeLength: 49, - text: 'a = 123'.repeat(7), - }, - ], - }) - assert.strictEqual(Object.keys(tracker.totalTokens).length, 1) - assert.strictEqual(Object.values(tracker.totalTokens)[0], 49) - }) - - it('Should skip when CodeWhisperer is editing', function () { - if (!tracker) { - assert.fail() - } - vsCodeState.isCodeWhispererEditing = true - tracker.countTotalTokens({ - reason: undefined, - document: createMockDocument(), - contentChanges: [ - { - range: new vscode.Range(0, 0, 0, 30), - rangeOffset: 0, - rangeLength: 30, - text: 'def twoSum(nums, target):\nfor', - }, - ], - }) - const startedSpy = sinon.spy(CodeWhispererCodeCoverageTracker.prototype, 'addTotalTokens') - assert.ok(!startedSpy.called) - }) - - it('Should not reduce tokens when delete', function () { - if (!tracker) { - assert.fail() - } - const doc = createMockDocument('import math', 'test.py', 'python') - tracker.countTotalTokens({ - reason: undefined, - document: doc, - contentChanges: [ - { - range: new vscode.Range(0, 0, 0, 1), - rangeOffset: 0, - rangeLength: 0, - text: 'a', - }, - ], - }) - tracker.countTotalTokens({ - reason: undefined, - document: doc, - contentChanges: [ - { - range: new vscode.Range(0, 0, 0, 1), - rangeOffset: 0, - rangeLength: 0, - text: 'b', - }, - ], - }) - assert.strictEqual(tracker?.totalTokens[doc.fileName], 2) - tracker.countTotalTokens({ - reason: undefined, - document: doc, - contentChanges: [ - { - range: new vscode.Range(0, 0, 0, 1), - rangeOffset: 1, - rangeLength: 1, - text: '', - }, - ], - }) - assert.strictEqual(tracker?.totalTokens[doc.fileName], 2) - }) - - it('Should add tokens when type', function () { - if (!tracker) { - assert.fail() - } - const doc = createMockDocument('import math', 'test.py', 'python') - tracker.countTotalTokens({ - reason: undefined, - document: doc, - contentChanges: [ - { - range: new vscode.Range(0, 0, 0, 1), - rangeOffset: 0, - rangeLength: 0, - text: 'a', - }, - ], - }) - assert.strictEqual(tracker?.totalTokens[doc.fileName], 1) - }) - - it('Should add tokens when hitting enter with indentation', function () { - if (!tracker) { - assert.fail() - } - const doc = createMockDocument('def h():', 'test.py', 'python') - tracker.countTotalTokens({ - reason: undefined, - document: doc, - contentChanges: [ - { - range: new vscode.Range(0, 0, 0, 8), - rangeOffset: 0, - rangeLength: 0, - text: '\n ', - }, - ], - }) - assert.strictEqual(tracker?.totalTokens[doc.fileName], 1) - }) - - it('Should add tokens when hitting enter with indentation in Windows', function () { - if (!tracker) { - assert.fail() - } - const doc = createMockDocument('def h():', 'test.py', 'python') - tracker.countTotalTokens({ - reason: undefined, - document: doc, - contentChanges: [ - { - range: new vscode.Range(0, 0, 0, 8), - rangeOffset: 0, - rangeLength: 0, - text: '\r\n ', - }, - ], - }) - assert.strictEqual(tracker?.totalTokens[doc.fileName], 1) - }) - - it('Should add tokens when hitting enter with indentation in Java', function () { - if (!tracker) { - assert.fail() - } - const doc = createMockDocument('class A() {', 'test.java', 'java') - tracker.countTotalTokens({ - reason: undefined, - document: doc, - contentChanges: [ - { - range: new vscode.Range(0, 0, 0, 11), - rangeOffset: 0, - rangeLength: 0, - text: '', - }, - { - range: new vscode.Range(0, 0, 0, 11), - rangeOffset: 0, - rangeLength: 0, - text: '\n\t\t', - }, - ], - }) - assert.strictEqual(tracker?.totalTokens[doc.fileName], 1) - }) - - it('Should add tokens when inserting closing brackets', function () { - if (!tracker) { - assert.fail() - } - const doc = createMockDocument('a=', 'test.py', 'python') - tracker.countTotalTokens({ - reason: undefined, - document: doc, - contentChanges: [ - { - range: new vscode.Range(0, 0, 0, 3), - rangeOffset: 0, - rangeLength: 0, - text: '[]', - }, - ], - }) - assert.strictEqual(tracker?.totalTokens[doc.fileName], 2) - }) - - it('Should add tokens when inserting closing brackets in Java', function () { - if (!tracker) { - assert.fail() - } - const doc = createMockDocument('class A ', 'test.java', 'java') - tracker.countTotalTokens({ - reason: undefined, - document: doc, - contentChanges: [ - { - range: new vscode.Range(0, 0, 0, 8), - rangeOffset: 0, - rangeLength: 0, - text: '{}', - }, - { - range: new vscode.Range(0, 0, 0, 8), - rangeOffset: 0, - rangeLength: 0, - text: '', - }, - ], - }) - assert.strictEqual(tracker?.totalTokens[doc.fileName], 2) - }) - }) - - describe('flush', function () { - let tracker: CodeWhispererCodeCoverageTracker | undefined - - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - tracker = CodeWhispererCodeCoverageTracker.getTracker(language) - if (tracker) { - sinon.stub(tracker, 'isActive').returns(true) - } - }) - - afterEach(function () { - sinon.restore() - CodeWhispererCodeCoverageTracker.instances.clear() - }) - - it('Should not send codecoverage telemetry if tracker is not active', function () { - if (!tracker) { - assert.fail() - } - sinon.restore() - sinon.stub(tracker, 'isActive').returns(false) - - tracker.addAcceptedTokens(`test.py`, { range: new vscode.Range(0, 0, 0, 7), text: `print()`, accepted: 7 }) - tracker.addTotalTokens(`test.py`, 100) - tracker.flush() - const data = globals.telemetry.logger.query({ - metricName: 'codewhisperer_codePercentage', - excludeKeys: ['awsAccount'], - }) - assert.strictEqual(data.length, 0) - }) - }) - - describe('emitCodeWhispererCodeContribution', function () { - let tracker: CodeWhispererCodeCoverageTracker | undefined - - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - tracker = CodeWhispererCodeCoverageTracker.getTracker(language) - if (tracker) { - sinon.stub(tracker, 'isActive').returns(true) - } - }) - - afterEach(function () { - sinon.restore() - CodeWhispererCodeCoverageTracker.instances.clear() - }) - - it('should emit correct code coverage telemetry in python file', async function () { - const tracker = CodeWhispererCodeCoverageTracker.getTracker(language) - - const assertTelemetry = assertTelemetryCurried('codewhisperer_codePercentage') - tracker?.incrementServiceInvocationCount() - tracker?.addAcceptedTokens(`test.py`, { range: new vscode.Range(0, 0, 0, 7), text: `print()`, accepted: 7 }) - tracker?.addTotalTokens(`test.py`, 100) - tracker?.emitCodeWhispererCodeContribution() - assertTelemetry({ - codewhispererTotalTokens: 100, - codewhispererLanguage: language, - codewhispererAcceptedTokens: 7, - codewhispererSuggestedTokens: 7, - codewhispererPercentage: 7, - successCount: 1, - }) - }) - - it('should emit correct code coverage telemetry when success count = 0', async function () { - const tracker = CodeWhispererCodeCoverageTracker.getTracker('java') - - const assertTelemetry = assertTelemetryCurried('codewhisperer_codePercentage') - tracker?.addAcceptedTokens(`test.java`, { - range: new vscode.Range(0, 0, 0, 18), - text: `public static main`, - accepted: 18, - }) - tracker?.incrementServiceInvocationCount() - tracker?.incrementServiceInvocationCount() - tracker?.addTotalTokens(`test.java`, 30) - tracker?.emitCodeWhispererCodeContribution() - assertTelemetry({ - codewhispererTotalTokens: 30, - codewhispererLanguage: 'java', - codewhispererAcceptedTokens: 18, - codewhispererSuggestedTokens: 18, - codewhispererPercentage: 60, - successCount: 2, - }) - }) - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts b/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts index ed17c181ee5..a43720c81be 100644 --- a/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts @@ -82,8 +82,6 @@ describe('codewhispererTracker', function () { describe('emitTelemetryOnSuggestion', function () { it('Should call recordCodewhispererUserModification with suggestion event', async function () { - const testStartUrl = 'testStartUrl' - sinon.stub(AuthUtil.instance, 'startUrl').value(testStartUrl) const suggestion = createAcceptedSuggestionEntry() const assertTelemetry = assertTelemetryCurried('codewhisperer_userModification') await CodeWhispererTracker.getTracker().emitTelemetryOnSuggestion(suggestion) @@ -95,7 +93,7 @@ describe('codewhispererTracker', function () { codewhispererModificationPercentage: 1, codewhispererCompletionType: 'Line', codewhispererLanguage: 'java', - credentialStartUrl: testStartUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, codewhispererCharactersAccepted: suggestion.originalString.length, codewhispererCharactersModified: 0, }) diff --git a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts index 500eaf23080..1795639e1e2 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts @@ -4,419 +4,389 @@ */ import assert from 'assert' -import { - AuthStates, - AuthUtil, - amazonQScopes, - codeWhispererChatScopes, - codeWhispererCoreScopes, -} from 'aws-core-vscode/codewhisperer' -import { - assertTelemetry, - getTestWindow, - SeverityLevel, - createBuilderIdProfile, - createSsoProfile, - createTestAuth, - captureEventNTimes, -} from 'aws-core-vscode/test' -import { Auth, Connection, isAnySsoConnection, isBuilderIdConnection } from 'aws-core-vscode/auth' -import { globals, vscodeComponent } from 'aws-core-vscode/shared' - -const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start' +import * as sinon from 'sinon' +import * as path from 'path' +import { AuthUtil, amazonQScopes } from 'aws-core-vscode/codewhisperer' +import { createTestAuthUtil, TestFolder } from 'aws-core-vscode/test' +import { constants, cache } from 'aws-core-vscode/auth' +import { auth2 } from 'aws-core-vscode/auth' +import { mementoUtils, fs } from 'aws-core-vscode/shared' describe('AuthUtil', async function () { - let auth: ReturnType - let authUtil: AuthUtil + let auth: any beforeEach(async function () { - auth = createTestAuth(globals.globalState) - authUtil = new AuthUtil(auth) + await createTestAuthUtil() + auth = AuthUtil.instance }) afterEach(async function () { - await auth.logout() + sinon.restore() }) - it('if there is no valid AwsBuilderID conn, it will create one and use it', async function () { - getTestWindow().onDidShowQuickPick(async (picker) => { - await picker.untilReady() - picker.acceptItem(picker.items[1]) + describe('Auth state', function () { + it('login with BuilderId', async function () { + await auth.login(constants.builderIdStartUrl, constants.builderIdRegion) + assert.ok(auth.isConnected()) + assert.ok(auth.isBuilderIdConnection()) }) - await authUtil.connectToAwsBuilderId() - const conn = authUtil.conn - assert.strictEqual(conn?.type, 'sso') - assert.strictEqual(conn.label, 'AWS Builder ID') - assert.deepStrictEqual(conn.scopes, amazonQScopes) - }) + it('login with IDC', async function () { + await auth.login('https://example.awsapps.com/start', 'us-east-1') + assert.ok(auth.isConnected()) + assert.ok(auth.isIdcConnection()) + }) - it('if there IS an existing AwsBuilderID conn, it will upgrade the scopes and use it', async function () { - const existingBuilderId = await auth.createConnection( - createBuilderIdProfile({ scopes: codeWhispererCoreScopes }) - ) - getTestWindow().onDidShowQuickPick(async (picker) => { - await picker.untilReady() - picker.acceptItem(picker.items[1]) + it('identifies internal users', async function () { + await auth.login(constants.internalStartUrl, 'us-east-1') + assert.ok(auth.isInternalAmazonUser()) }) - await authUtil.connectToAwsBuilderId() + it('identifies SSO session', function () { + ;(auth as any).session = { loginType: auth2.LoginTypes.SSO } + assert.strictEqual(auth.isSsoSession(), true) + }) - const conn = authUtil.conn - assert.strictEqual(conn?.type, 'sso') - assert.strictEqual(conn.id, existingBuilderId.id) - assert.deepStrictEqual(conn.scopes, amazonQScopes) + it('identifies non-SSO session', function () { + ;(auth as any).session = { loginType: auth2.LoginTypes.IAM } + assert.strictEqual(auth.isSsoSession(), false) + }) }) - it('if there is no valid enterprise SSO conn, will create and use one', async function () { - getTestWindow().onDidShowQuickPick(async (picker) => { - await picker.untilReady() - picker.acceptItem(picker.items[1]) + describe('Token management', function () { + it('can get token when connected with SSO', async function () { + await auth.login(constants.builderIdStartUrl, constants.builderIdRegion) + const token = await auth.getToken() + assert.ok(token) }) - await authUtil.connectToEnterpriseSso(enterpriseSsoStartUrl, 'us-east-1') - const conn = authUtil.conn - assert.strictEqual(conn?.type, 'sso') - assert.strictEqual(conn.label, 'IAM Identity Center (enterprise)') + it('throws when getting token without SSO connection', async function () { + sinon.stub(AuthUtil.instance, 'isSsoSession').returns(false) + await assert.rejects(async () => await auth.getToken()) + }) }) - it('should add scopes + connect to existing IAM Identity Center connection', async function () { - getTestWindow().onDidShowMessage(async (message) => { - assert.ok(message.modal) - message.selectItem('Proceed') + describe('getTelemetryMetadata', function () { + it('returns valid metadata for BuilderId connection', async function () { + await auth.login(constants.builderIdStartUrl, constants.builderIdRegion) + const metadata = await auth.getTelemetryMetadata() + assert.strictEqual(metadata.credentialSourceId, 'awsId') + assert.strictEqual(metadata.credentialStartUrl, constants.builderIdStartUrl) }) - const randomScope = 'my:random:scope' - const ssoConn = await auth.createInvalidSsoConnection( - createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: [randomScope] }) - ) - // Method under test - await authUtil.connectToEnterpriseSso(ssoConn.startUrl, 'us-east-1') + it('returns valid metadata for IDC connection', async function () { + await auth.login('https://example.awsapps.com/start', 'us-east-1') + const metadata = await auth.getTelemetryMetadata() + assert.strictEqual(metadata.credentialSourceId, 'iamIdentityCenter') + assert.strictEqual(metadata.credentialStartUrl, 'https://example.awsapps.com/start') + }) - const cwConn = authUtil.conn - assert.strictEqual(cwConn?.type, 'sso') - assert.strictEqual(cwConn.label, 'IAM Identity Center (enterprise)') - assert.deepStrictEqual(cwConn.scopes, [randomScope, ...amazonQScopes]) + it('returns undefined metadata when not connected', async function () { + await auth.logout() + const metadata = await auth.getTelemetryMetadata() + assert.strictEqual(metadata.id, 'undefined') + }) }) - it('reauthenticates an existing BUT invalid Amazon Q IAM Identity Center connection', async function () { - const ssoConn = await auth.createInvalidSsoConnection( - createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: amazonQScopes }) - ) - await auth.refreshConnectionState(ssoConn) - assert.strictEqual(auth.getConnectionState(ssoConn), 'invalid') - - // Method under test - await authUtil.connectToEnterpriseSso(ssoConn.startUrl, 'us-east-1') - - const cwConn = authUtil.conn - assert.strictEqual(cwConn?.type, 'sso') - assert.strictEqual(cwConn.id, ssoConn.id) - assert.deepStrictEqual(cwConn.scopes, amazonQScopes) - assert.strictEqual(auth.getConnectionState(cwConn), 'valid') - }) + describe('getAuthFormIds', function () { + it('returns empty array when not connected', async function () { + await auth.logout() + const forms = await auth.getAuthFormIds() + assert.deepStrictEqual(forms, []) + }) - it('should show reauthenticate prompt', async function () { - getTestWindow().onDidShowMessage((m) => { - if (m.severity === SeverityLevel.Information) { - m.close() - } + it('returns BuilderId forms when using BuilderId', async function () { + await auth.login(constants.builderIdStartUrl, constants.builderIdRegion) + const forms = await auth.getAuthFormIds() + assert.deepStrictEqual(forms, ['builderIdCodeWhisperer']) }) - await authUtil.showReauthenticatePrompt() + it('returns IDC forms when using IDC without SSO account access', async function () { + const session = (auth as any).session + sinon.stub(session, 'getProfile').resolves({ + ssoSession: { + settings: { + sso_registration_scopes: ['codewhisperer:*'], + }, + }, + }) - const warningMessage = getTestWindow().shownMessages.filter((m) => m.severity === SeverityLevel.Information) - assert.strictEqual(warningMessage.length, 1) - assert.strictEqual(warningMessage[0].message, `Your Amazon Q connection has expired. Please re-authenticate.`) - warningMessage[0].close() - assertTelemetry('toolkit_showNotification', { - id: 'codeWhispererConnectionExpired', - result: 'Succeeded', - source: vscodeComponent, + await auth.login('https://example.awsapps.com/start', 'us-east-1') + const forms = await auth.getAuthFormIds() + assert.deepStrictEqual(forms, ['identityCenterCodeWhisperer']) }) - assertTelemetry('toolkit_invokeAction', { - id: 'codeWhispererConnectionExpired', - action: 'dismiss', - result: 'Succeeded', - source: vscodeComponent, + + it('returns IDC forms with explorer when using IDC with SSO account access', async function () { + const session = (auth as any).session + sinon.stub(session, 'getProfile').resolves({ + ssoSession: { + settings: { + sso_registration_scopes: ['codewhisperer:*', 'sso:account:access'], + }, + }, + }) + + await auth.login('https://example.awsapps.com/start', 'us-east-1') + const forms = await auth.getAuthFormIds() + assert.deepStrictEqual(forms.sort(), ['identityCenterCodeWhisperer', 'identityCenterExplorer'].sort()) + }) + + it('returns credentials form for IAM credentials', async function () { + sinon.stub(auth, 'isSsoSession').returns(false) + sinon.stub(auth, 'isConnected').returns(true) + + const forms = await auth.getAuthFormIds() + assert.deepStrictEqual(forms, ['credentials']) }) }) - it('reauthenticate prompt reauthenticates invalid connection', async function () { - const conn = await auth.createInvalidSsoConnection( - createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: codeWhispererChatScopes }) - ) - await auth.useConnection(conn) - getTestWindow().onDidShowMessage((m) => { - m.selectItem('Re-authenticate') + describe('cacheChangedHandler', function () { + it('calls logout when event is delete', async function () { + const logoutSpy = sinon.spy(auth, 'logout') + + await (auth as any).cacheChangedHandler('delete') + + assert.ok(logoutSpy.calledOnce) }) - assert.strictEqual(auth.getConnectionState(conn), 'invalid') + it('calls restore when event is create', async function () { + const restoreSpy = sinon.spy(auth, 'restore') - await authUtil.showReauthenticatePrompt() + await (auth as any).cacheChangedHandler('create') - assert.strictEqual(authUtil.conn?.type, 'sso') - assert.strictEqual(auth.getConnectionState(conn), 'valid') - assertTelemetry('toolkit_showNotification', { - id: 'codeWhispererConnectionExpired', - result: 'Succeeded', - source: vscodeComponent, + assert.ok(restoreSpy.calledOnce) }) - assertTelemetry('toolkit_invokeAction', { - id: 'codeWhispererConnectionExpired', - action: 'connect', - result: 'Succeeded', - source: vscodeComponent, + + it('does nothing for other events', async function () { + const logoutSpy = sinon.spy(auth, 'logout') + const restoreSpy = sinon.spy(auth, 'restore') + + await (auth as any).cacheChangedHandler('unknown') + + assert.ok(logoutSpy.notCalled) + assert.ok(restoreSpy.notCalled) }) }) - it('reauthenticates Builder ID connection that already has all scopes', async function () { - const conn = await auth.createInvalidSsoConnection(createBuilderIdProfile({ scopes: amazonQScopes })) - await auth.useConnection(conn) + describe('stateChangeHandler', function () { + let mockLspAuth: any + let regionProfileManager: any - // method under test - await authUtil.reauthenticate() + beforeEach(function () { + mockLspAuth = (auth as any).lspAuth + regionProfileManager = (auth as any).regionProfileManager + }) - assert.strictEqual(authUtil.conn?.type, 'sso') - assert.deepStrictEqual(authUtil.conn?.scopes, amazonQScopes) - assert.strictEqual(auth.getConnectionState(conn), 'valid') - }) + it('updates bearer token when state is refreshed', async function () { + await auth.login(constants.builderIdStartUrl, 'us-east-1') - it('reauthenticates IdC connection that already has all scopes', async function () { - const conn = await auth.createInvalidSsoConnection( - createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: codeWhispererCoreScopes }) - ) - await auth.useConnection(conn) + await (auth as any).stateChangeHandler({ state: 'refreshed' }) - // method under test - await authUtil.reauthenticate() + assert.ok(mockLspAuth.updateBearerToken.called) + assert.strictEqual(mockLspAuth.updateBearerToken.firstCall.args[0].data, 'fake-data') + }) - assert.strictEqual(authUtil.conn?.type, 'sso') - assert.deepStrictEqual(authUtil.conn?.scopes, amazonQScopes) - assert.strictEqual(auth.getConnectionState(conn), 'valid') - }) + it('cleans up when connection expires', async function () { + await auth.login(constants.builderIdStartUrl, 'us-east-1') - it('reauthenticate adds missing Builder ID scopes', async function () { - const conn = await auth.createInvalidSsoConnection(createBuilderIdProfile({ scopes: codeWhispererCoreScopes })) - await auth.useConnection(conn) + await (auth as any).stateChangeHandler({ state: 'expired' }) - // method under test - await authUtil.reauthenticate() + assert.ok(mockLspAuth.deleteBearerToken.called) + }) - assert.strictEqual(authUtil.conn?.type, 'sso') - assert.deepStrictEqual(authUtil.conn?.scopes, amazonQScopes) - assert.strictEqual(auth.getConnectionState(conn), 'valid') - }) + it('deletes bearer token when disconnected', async function () { + await (auth as any).stateChangeHandler({ state: 'notConnected' }) - it('reauthenticate adds missing Amazon Q IdC scopes', async function () { - const conn = await auth.createInvalidSsoConnection( - createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: codeWhispererCoreScopes }) - ) - await auth.useConnection(conn) + assert.ok(mockLspAuth.deleteBearerToken.called) + }) - // method under test - await authUtil.reauthenticate() + it('updates bearer token and restores profile on reconnection', async function () { + const restoreProfileSelectionSpy = sinon.spy(regionProfileManager, 'restoreProfileSelection') - assert.strictEqual(authUtil.conn?.type, 'sso') - assert.deepStrictEqual(authUtil.conn?.scopes, amazonQScopes) - assert.strictEqual(auth.getConnectionState(conn), 'valid') - }) + await auth.login('https://example.awsapps.com/start', 'us-east-1') - it('CodeWhisperer uses fallback connection when switching to an unsupported connection', async function () { - const supportedConn = await auth.createConnection(createBuilderIdProfile({ scopes: codeWhispererChatScopes })) - const unsupportedConn = await auth.createConnection(createSsoProfile()) - - await auth.useConnection(supportedConn) - assert.ok(authUtil.isConnected()) - assert.strictEqual(auth.activeConnection?.id, authUtil.conn?.id) - - // Switch to unsupported connection - const cwAuthUpdatedConnection = captureEventNTimes(authUtil.secondaryAuth.onDidChangeActiveConnection, 2) - await auth.useConnection(unsupportedConn) - // - This is triggered when the main Auth connection is switched - // - This is triggered by registerAuthListener() when it saves the previous active connection as a fallback. - await cwAuthUpdatedConnection - - // TODO in a refactor see if we can simplify multiple multiple triggers on the same event. - assert.ok(authUtil.isConnected()) - assert.ok(authUtil.isUsingSavedConnection) - assert.notStrictEqual(auth.activeConnection?.id, authUtil.conn?.id) - assert.strictEqual(authUtil.conn?.type, 'sso') - assert.deepStrictEqual(authUtil.conn?.scopes, codeWhispererChatScopes) - }) + await (auth as any).stateChangeHandler({ state: 'connected' }) - it('does not prompt to sign out of duplicate builder ID connections', async function () { - await authUtil.connectToAwsBuilderId() - await authUtil.connectToAwsBuilderId() - assert.ok(authUtil.isConnected()) + assert.ok(mockLspAuth.updateBearerToken.called) + assert.ok(restoreProfileSelectionSpy.called) + }) - const ssoConnectionIds = new Set(auth.activeConnectionEvents.emits.filter(isAnySsoConnection).map((c) => c.id)) - assert.strictEqual(ssoConnectionIds.size, 1, 'Expected exactly 1 unique SSO connection id') - assert.strictEqual((await auth.listConnections()).filter(isAnySsoConnection).length, 1) - }) + it('clears region profile cache and invalidates profile on IDC connection expiration', async function () { + const invalidateProfileSpy = sinon.spy(regionProfileManager, 'invalidateProfile') + const clearCacheSpy = sinon.spy(regionProfileManager, 'clearCache') - it('automatically upgrades connections if they do not have the required scopes', async function () { - const upgradeableConn = await auth.createConnection(createBuilderIdProfile()) - await auth.useConnection(upgradeableConn) - assert.strictEqual(authUtil.isConnected(), false) - - await authUtil.connectToAwsBuilderId() - assert.ok(authUtil.isConnected()) - assert.ok(authUtil.isConnectionValid()) - assert.ok(isBuilderIdConnection(authUtil.conn)) - assert.strictEqual(authUtil.conn?.id, upgradeableConn.id) - assert.strictEqual(authUtil.conn.startUrl, upgradeableConn.startUrl) - assert.strictEqual(authUtil.conn.ssoRegion, upgradeableConn.ssoRegion) - assert.deepStrictEqual(authUtil.conn.scopes, amazonQScopes) - assert.strictEqual((await auth.listConnections()).filter(isAnySsoConnection).length, 1) - }) + await auth.login('https://example.awsapps.com/start', 'us-east-1') - it('test reformatStartUrl should remove trailing slash and hash', function () { - const expected = 'https://view.awsapps.com/start' - assert.strictEqual(authUtil.reformatStartUrl(expected + '/'), expected) - assert.strictEqual(authUtil.reformatStartUrl(undefined), undefined) - assert.strictEqual(authUtil.reformatStartUrl(expected + '/#'), expected) - assert.strictEqual(authUtil.reformatStartUrl(expected + '#/'), expected) - assert.strictEqual(authUtil.reformatStartUrl(expected + '/#/'), expected) - assert.strictEqual(authUtil.reformatStartUrl(expected + '####'), expected) - }) + await (auth as any).stateChangeHandler({ state: 'expired' }) - it(`clearExtraConnections()`, async function () { - const conn1 = await auth.createConnection(createBuilderIdProfile()) - const conn2 = await auth.createConnection(createSsoProfile({ startUrl: enterpriseSsoStartUrl })) - const conn3 = await auth.createConnection(createSsoProfile({ startUrl: enterpriseSsoStartUrl + 1 })) - // validate listConnections shows all connections - assert.deepStrictEqual( - (await authUtil.auth.listConnections()).map((conn) => conn.id).sort((a, b) => a.localeCompare(b)), - [conn1, conn2, conn3].map((conn) => conn.id).sort((a, b) => a.localeCompare(b)) - ) - await authUtil.secondaryAuth.useNewConnection(conn3) - - await authUtil.clearExtraConnections() // method under test - - // Only the conn that AuthUtil is using is remaining - assert.deepStrictEqual( - (await authUtil.auth.listConnections()).map((conn) => conn.id), - [conn3.id] - ) + assert.ok(invalidateProfileSpy.called) + assert.ok(clearCacheSpy.called) + }) }) -}) -describe('getChatAuthState()', function () { - let auth: ReturnType - let authUtil: AuthUtil - let laterDate: Date + describe('migrateSsoConnectionToLsp', function () { + let mockLspAuth: any + let memento: any + let cacheDir: string + let fromRegistrationFile: string + let fromTokenFile: string + + const validProfile = { + type: 'sso', + startUrl: 'https://test2.com', + ssoRegion: 'us-east-1', + scopes: amazonQScopes, + metadata: { + connectionState: 'valid', + }, + } + + beforeEach(async function () { + memento = { + get: sinon.stub(), + update: sinon.stub().resolves(), + } + cacheDir = (await TestFolder.create()).path + + sinon.stub(mementoUtils, 'getEnvironmentSpecificMemento').returns(memento) + sinon.stub(cache, 'getCacheDir').returns(cacheDir) - beforeEach(async function () { - auth = createTestAuth(globals.globalState) - authUtil = new AuthUtil(auth) + mockLspAuth = (auth as any).lspAuth + mockLspAuth.getSsoToken.resolves(undefined) - laterDate = new Date(Date.now() + 10_000_000) - }) + fromTokenFile = cache.getTokenCacheFile(cacheDir, 'profile1') + const registrationKey = { + startUrl: validProfile.startUrl, + region: validProfile.ssoRegion, + scopes: amazonQScopes, + } + fromRegistrationFile = cache.getRegistrationCacheFile(cacheDir, registrationKey) - afterEach(async function () { - await auth.logout() - }) + const registrationData = { test: 'registration' } + const tokenData = { test: 'token' } - it('indicates nothing connected when no auth connection exists', async function () { - const result = await authUtil.getChatAuthState() - assert.deepStrictEqual(result, { - codewhispererChat: AuthStates.disconnected, - codewhispererCore: AuthStates.disconnected, - amazonQ: AuthStates.disconnected, + await fs.writeFile(fromRegistrationFile, JSON.stringify(registrationData)) + await fs.writeFile(fromTokenFile, JSON.stringify(tokenData)) }) - }) - /** Affects {@link Auth.refreshConnectionState} */ - function createToken(conn: Connection) { - auth.getTestTokenProvider(conn).getToken.resolves({ accessToken: 'myAccessToken', expiresAt: laterDate }) - } - - describe('Builder ID', function () { - it('indicates only CodeWhisperer core is connected when only CW core scopes are set', async function () { - const conn = await auth.createConnection(createBuilderIdProfile({ scopes: codeWhispererCoreScopes })) - createToken(conn) - await auth.useConnection(conn) - - const result = await authUtil.getChatAuthState() - assert.deepStrictEqual(result, { - codewhispererCore: AuthStates.connected, - codewhispererChat: AuthStates.expired, - amazonQ: AuthStates.expired, - }) + afterEach(async function () { + sinon.restore() }) - it('indicates all SUPPORTED features connected when all scopes are set', async function () { - const conn = await auth.createConnection(createBuilderIdProfile({ scopes: amazonQScopes })) - createToken(conn) - await auth.useConnection(conn) + it('skips migration if LSP token exists', async function () { + memento.get.returns({ profile1: validProfile }) + mockLspAuth.getSsoToken.resolves({ token: 'valid-token' }) - const result = await authUtil.getChatAuthState() - assert.deepStrictEqual(result, { - codewhispererCore: AuthStates.connected, - codewhispererChat: AuthStates.connected, - amazonQ: AuthStates.connected, - }) + await auth.migrateSsoConnectionToLsp('test-client') + + assert.ok(memento.update.calledWith('auth.profiles', undefined)) + assert.ok(!auth.session.updateProfile?.called) }) - it('indicates all SUPPORTED features expired when connection is invalid', async function () { - const conn = await auth.createInvalidSsoConnection( - createBuilderIdProfile({ scopes: codeWhispererChatScopes }) - ) - await auth.useConnection(conn) + it('proceeds with migration if LSP token check throws', async function () { + memento.get.returns({ profile1: validProfile }) + mockLspAuth.getSsoToken.rejects(new Error('Token check failed')) + const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves() - const result = await authUtil.getChatAuthState() - assert.deepStrictEqual(result, { - codewhispererCore: AuthStates.expired, - codewhispererChat: AuthStates.expired, - amazonQ: AuthStates.expired, - }) + await auth.migrateSsoConnectionToLsp('test-client') + + assert.ok(updateProfileStub.calledOnce) + assert.ok(memento.update.calledWith('auth.profiles', undefined)) }) - }) - describe('Identity Center', function () { - it('indicates only CW core is connected when only CW core scopes are set', async function () { - const conn = await auth.createConnection( - createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: codeWhispererCoreScopes }) - ) - createToken(conn) - await auth.useConnection(conn) - - const result = await authUtil.getChatAuthState() - assert.deepStrictEqual(result, { - codewhispererCore: AuthStates.pendingProfileSelection, - codewhispererChat: AuthStates.expired, - amazonQ: AuthStates.expired, - }) + it('migrates valid SSO connection', async function () { + memento.get.returns({ profile1: validProfile }) + + const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves() + + await auth.migrateSsoConnectionToLsp('test-client') + + assert.ok(updateProfileStub.calledOnce) + assert.ok(memento.update.calledWith('auth.profiles', undefined)) + + const files = await fs.readdir(cacheDir) + assert.strictEqual(files.length, 2) // Should have both the token and registration file + + // Verify file contents were preserved + const newFiles = files.map((f) => path.join(cacheDir, f[0])) + for (const file of newFiles) { + const content = await fs.readFileText(file) + const parsed = JSON.parse(content) + assert.ok(parsed.test === 'registration' || parsed.test === 'token') + } }) - it('indicates all features connected when all scopes are set', async function () { - const conn = await auth.createConnection( - createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: amazonQScopes }) - ) - createToken(conn) - await auth.useConnection(conn) - - const result = await authUtil.getChatAuthState() - assert.deepStrictEqual(result, { - codewhispererCore: AuthStates.pendingProfileSelection, - codewhispererChat: AuthStates.pendingProfileSelection, - amazonQ: AuthStates.pendingProfileSelection, - }) + it('does not migrate if no matching SSO profile exists', async function () { + const mockProfiles = { + 'test-profile': { + type: 'iam', + startUrl: 'https://test.com', + ssoRegion: 'us-east-1', + }, + } + memento.get.returns(mockProfiles) + + await auth.migrateSsoConnectionToLsp('test-client') + + // Assert that the file names have not updated + const files = await fs.readdir(cacheDir) + assert.ok(files.length === 2) + assert.ok(await fs.exists(fromRegistrationFile)) + assert.ok(await fs.exists(fromTokenFile)) + assert.ok(!memento.update.called) }) - it('indicates all features expired when connection is invalid', async function () { - const conn = await auth.createInvalidSsoConnection( - createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: amazonQScopes }) - ) - await auth.useConnection(conn) + it('migrates only profile with matching scopes', async function () { + const mockProfiles = { + profile1: validProfile, + profile2: { + type: 'sso', + startUrl: 'https://test.com', + ssoRegion: 'us-east-1', + scopes: ['different:scope'], + metadata: { + connectionState: 'valid', + }, + }, + } + memento.get.returns(mockProfiles) - const result = await authUtil.getChatAuthState() - assert.deepStrictEqual(result, { - codewhispererCore: AuthStates.expired, - codewhispererChat: AuthStates.expired, - amazonQ: AuthStates.expired, + const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves() + + await auth.migrateSsoConnectionToLsp('test-client') + + assert.ok(updateProfileStub.calledOnce) + assert.ok(memento.update.calledWith('auth.profiles', undefined)) + assert.deepStrictEqual(updateProfileStub.firstCall.args[0], { + startUrl: validProfile.startUrl, + region: validProfile.ssoRegion, + scopes: validProfile.scopes, }) }) + + it('uses valid connection state when multiple profiles exist', async function () { + const mockProfiles = { + profile2: { + ...validProfile, + metadata: { + connectionState: 'invalid', + }, + }, + profile1: validProfile, + } + memento.get.returns(mockProfiles) + + const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves() + + await auth.migrateSsoConnectionToLsp('test-client') + + assert.ok( + updateProfileStub.calledWith({ + startUrl: validProfile.startUrl, + region: validProfile.ssoRegion, + scopes: validProfile.scopes, + }) + ) + }) }) }) diff --git a/packages/amazonq/test/unit/codewhisperer/util/bm25.test.ts b/packages/amazonq/test/unit/codewhisperer/util/bm25.test.ts deleted file mode 100644 index 0a3c4b17d60..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/util/bm25.test.ts +++ /dev/null @@ -1,117 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import { BM25Okapi } from 'aws-core-vscode/codewhisperer' - -describe('bm25', function () { - it('simple case 1', function () { - const query = 'windy London' - const corpus = ['Hello there good man!', 'It is quite windy in London', 'How is the weather today?'] - - const sut = new BM25Okapi(corpus) - const actual = sut.score(query) - - assert.deepStrictEqual(actual, [ - { - content: 'Hello there good man!', - index: 0, - score: 0, - }, - { - content: 'It is quite windy in London', - index: 1, - score: 0.937294722506405, - }, - { - content: 'How is the weather today?', - index: 2, - score: 0, - }, - ]) - - assert.deepStrictEqual(sut.topN(query, 1), [ - { - content: 'It is quite windy in London', - index: 1, - score: 0.937294722506405, - }, - ]) - }) - - it('simple case 2', function () { - const query = 'codewhisperer is a machine learning powered code generator' - const corpus = [ - 'codewhisperer goes GA at April 2023', - 'machine learning tool is the trending topic!!! :)', - 'codewhisperer is good =))))', - 'codewhisperer vs. copilot, which code generator better?', - 'copilot is a AI code generator too', - 'it is so amazing!!', - ] - - const sut = new BM25Okapi(corpus) - const actual = sut.score(query) - - assert.deepStrictEqual(actual, [ - { - content: 'codewhisperer goes GA at April 2023', - index: 0, - score: 0, - }, - { - content: 'machine learning tool is the trending topic!!! :)', - index: 1, - score: 2.597224531416621, - }, - { - content: 'codewhisperer is good =))))', - index: 2, - score: 0.3471790843435529, - }, - { - content: 'codewhisperer vs. copilot, which code generator better?', - index: 3, - score: 1.063018436525109, - }, - { - content: 'copilot is a AI code generator too', - index: 4, - score: 2.485359418462239, - }, - { - content: 'it is so amazing!!', - index: 5, - score: 0.3154033715392277, - }, - ]) - - assert.deepStrictEqual(sut.topN(query, 1), [ - { - content: 'machine learning tool is the trending topic!!! :)', - index: 1, - score: 2.597224531416621, - }, - ]) - - assert.deepStrictEqual(sut.topN(query, 3), [ - { - content: 'machine learning tool is the trending topic!!! :)', - index: 1, - score: 2.597224531416621, - }, - { - content: 'copilot is a AI code generator too', - index: 4, - score: 2.485359418462239, - }, - { - content: 'codewhisperer vs. copilot, which code generator better?', - index: 3, - score: 1.063018436525109, - }, - ]) - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/util/closingBracketUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/closingBracketUtil.test.ts deleted file mode 100644 index bfdf9dc3d29..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/util/closingBracketUtil.test.ts +++ /dev/null @@ -1,389 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import assert from 'assert' -import { handleExtraBrackets } from 'aws-core-vscode/codewhisperer' -import { toTextEditor } from 'aws-core-vscode/test' - -describe('closingBracketUtil', function () { - /** - * leftContext + recommendation + rightContext - * startStart start end endEnd - */ - describe('handleExtraBrackets', function () { - async function assertClosingSymbolsHandler( - leftContext: string, - rightContext: string, - recommendation: string, - expected: string - ) { - const editor = await toTextEditor(leftContext + recommendation + rightContext, 'test.txt') - const document = editor.document - - const startStart = document.positionAt(0) - const endEnd = document.positionAt(editor.document.getText().length) - const start = document.positionAt(leftContext.length) - const end = document.positionAt(leftContext.length + recommendation.length) - - const left = document.getText(new vscode.Range(startStart, start)) - const right = document.getText(new vscode.Range(end, endEnd)) - const reco = document.getText(new vscode.Range(start, end)) - - assert.strictEqual(left, leftContext) - assert.strictEqual(right, rightContext) - assert.strictEqual(reco, recommendation) - - await handleExtraBrackets(editor, end, start) - - assert.strictEqual(editor.document.getText(), expected) - } - - it('should remove extra closing symbol', async function () { - /** - * public static void mergeSort(int[|] nums) { - * mergeSort(nums, 0, nums.length - 1); - * }|]) - */ - await assertClosingSymbolsHandler( - String.raw`public static void mergeSort(int[`, - String.raw`])`, - String.raw`] nums) { - mergeSort(nums, 0, nums.length - 1); -}`, - String.raw`public static void mergeSort(int[] nums) { - mergeSort(nums, 0, nums.length - 1); -}` - ) - - /** - * fun genericFunction<|T>(value: T): T { - * return value - * }|> - */ - await assertClosingSymbolsHandler( - String.raw`fun genericFunction<`, - String.raw`>`, - String.raw`T>(value: T): T { - return value -}`, - String.raw`fun genericFunction(value: T): T { - return value -}` - ) - - /** - * function getProperty(obj: T, key: K) {|> - */ - await assertClosingSymbolsHandler( - String.raw`function getProperty`, - String.raw`K extends keyof T>(obj: T, key: K) {`, - String.raw`function getProperty(obj: T, key: K) {` - ) - - /** - * public class Main { - * public static void main(|args: String[]) { - * System.out.println("Hello World"); - * }|) - * } - */ - await assertClosingSymbolsHandler( - String.raw`public class Main { - public static void main(`, - String.raw`) -}`, - String.raw`args: String[]) { - System.out.println("Hello World"); - }`, - String.raw`public class Main { - public static void main(args: String[]) { - System.out.println("Hello World"); - } -}` - ) - - /** - * function add2Numbers(a: number: b: number) { - * return a + b - * }) - */ - await assertClosingSymbolsHandler( - 'function add2Numbers(', - ')', - 'a: number, b: number) {\n return a + b\n}', - `function add2Numbers(a: number, b: number) {\n return a + b\n}` - ) - - /** - * function sum(a: number, b: number, c: number) { - * return a + b + c - * }) - */ - await assertClosingSymbolsHandler( - 'function sum(a: number, b: number, ', - ')', - 'c: number) {\n return a + b + c\n}', - `function sum(a: number, b: number, c: number) {\n return a + b + c\n}` - ) - - /** - * const aString = "hello world";" - */ - await assertClosingSymbolsHandler( - 'const aString = "', - '"', - 'hello world";', - `const aString = "hello world";` - ) - - /** - * { - * "userName": "john", - * "department": "codewhisperer"", - * } - */ - await assertClosingSymbolsHandler( - '{\n\t"userName": "john",\n\t"', - '"\n}', - 'department": "codewhisperer",', - '{\n\t"userName": "john",\n\t"department": "codewhisperer",\n}' - ) - - /** - * const someArray = [|"element1", "element2"];|] - */ - await assertClosingSymbolsHandler( - 'const anArray = [', - ']', - '"element1", "element2"];', - `const anArray = ["element1", "element2"];` - ) - - /** - * export const launchTemplates: { [key: string]: AmazonEC2.LaunchTemplate } = { - * lt1: { launchTemplateId: "lt-1", launchTemplateName: "foo" }, - * lt2: { launchTemplateId: "lt-2345", launchTemplateName: "bar" }, - * lt3: { |launchTemplateId: "lt-678919", launchTemplateName: "foobar" },| - * }; - */ - await assertClosingSymbolsHandler( - String.raw`export const launchTemplates: { [key: string]: AmazonEC2.LaunchTemplate } = { - lt1: { launchTemplateId: "lt-1", launchTemplateName: "foo" }, - lt2: { launchTemplateId: "lt-2345", launchTemplateName: "bar" }, - lt3: { `, - String.raw` - };`, - String.raw`launchTemplateId: "lt-678919", launchTemplateName: "foobar" },`, - String.raw`export const launchTemplates: { [key: string]: AmazonEC2.LaunchTemplate } = { - lt1: { launchTemplateId: "lt-1", launchTemplateName: "foo" }, - lt2: { launchTemplateId: "lt-2345", launchTemplateName: "bar" }, - lt3: { launchTemplateId: "lt-678919", launchTemplateName: "foobar" }, - };` - ) - - /** - * genericFunction<|T>|> () { - * if (T isInstanceOf string) { - * console.log(T) - * } else { - * // Do nothing - * } - * } - */ - await assertClosingSymbolsHandler( - String.raw`genericFunction<`, - String.raw`> () { - if (T isInstanceOf string) { - console.log(T) - } else { - // Do nothing - } -}`, - 'T>', - String.raw`genericFunction () { - if (T isInstanceOf string) { - console.log(T) - } else { - // Do nothing - } -}` - ) - - /** - * const rawStr = "|Foo";|" - * const anotherStr = "Bar" - */ - await assertClosingSymbolsHandler( - 'const rawStr = "', - '\nconst anotherStr = "Bar";', - 'Foo";', - String.raw`const rawStr = "Foo"; -const anotherStr = "Bar";` - ) - }) - - it('should not remove extra closing symbol', async function () { - /** - * describe('Foo', () => { - * describe('Bar', function () => { - * it('Boo', |() => { - * expect(true).toBe(true) - * }|) - * }) - * }) - */ - await assertClosingSymbolsHandler( - String.raw`describe('Foo', () => { - describe('Bar', function () { - it('Boo', `, - String.raw`) - }) -})`, - String.raw`() => { - expect(true).toBe(true) - }`, - String.raw`describe('Foo', () => { - describe('Bar', function () { - it('Boo', () => { - expect(true).toBe(true) - }) - }) -})` - ) - - /** - * function add2Numbers(|a: nuumber, b: number) { - * return a + b; - * }| - */ - await assertClosingSymbolsHandler( - 'function add2Numbers(', - '', - 'a: number, b: number) {\n return a + b;\n}', - `function add2Numbers(a: number, b: number) {\n return a + b;\n}` - ) - - /** - * export const launchTemplates: { [key: string]: AmazonEC2.LaunchTemplate } = { - * lt1: { launchTemplateId: "lt-1", launchTemplateName: "foo" }, - * lt2: { launchTemplateId: "lt-2345", launchTemplateName: "bar" }, - * lt3: |{ launchTemplateId: "lt-3456", launchTemplateName: "baz" },| - * } - */ - await assertClosingSymbolsHandler( - 'export const launchTemplates: { [key: string]: AmazonEC2.LaunchTemplate } = {\n lt1: { launchTemplateId: "lt-1", launchTemplateName: "foo" },\n lt2: { launchTemplateId: "lt-2345", launchTemplateName: "bar" },\n lt3: ', - '\n};', - '{ launchTemplateId: "lt-3456", launchTemplateName: "baz" },', - `export const launchTemplates: { [key: string]: AmazonEC2.LaunchTemplate } = {\n lt1: { launchTemplateId: "lt-1", launchTemplateName: "foo" },\n lt2: { launchTemplateId: "lt-2345", launchTemplateName: "bar" },\n lt3: { launchTemplateId: "lt-3456", launchTemplateName: "baz" },\n};` - ) - - /** - * export const launchTemplates: { [key: string]: AmazonEC2.LaunchTemplate } = { - * lt1: { launchTemplateId: "lt-1", launchTemplateName: "foo" }, - * lt2: { launchTemplateId: "lt-2345", launchTemplateName: "bar" }, - * |lt3: { launchTemplateId: "lt-3456", launchTemplateName: "baz" },| - * } - */ - await assertClosingSymbolsHandler( - 'export const launchTemplates: { [key: string]: AmazonEC2.LaunchTemplate } = {\n lt1: { launchTemplateId: "lt-1", launchTemplateName: "foo" },\n lt2: { launchTemplateId: "lt-2345", launchTemplateName: "bar" },\n ', - '\n};', - 'lt3: { launchTemplateId: "lt-3456", launchTemplateName: "baz" },', - 'export const launchTemplates: { [key: string]: AmazonEC2.LaunchTemplate } = {\n lt1: { launchTemplateId: "lt-1", launchTemplateName: "foo" },\n lt2: { launchTemplateId: "lt-2345", launchTemplateName: "bar" },\n lt3: { launchTemplateId: "lt-3456", launchTemplateName: "baz" },\n};' - ) - - /** - * const aString = "|hello world";| - */ - await assertClosingSymbolsHandler( - 'const aString = "', - '', - 'hello world";', - 'const aString = "hello world";' - ) - - /** genericFunction<|T> ()|> { - * if (T isInstanceOf string) { - * console.log(T) - * } else { - * // Do nothing - * } - * } - */ - await assertClosingSymbolsHandler( - 'genericFunction<', - String.raw` { - if (T isInstanceOf string) { - console.log(T) - } else { - // Do nothing - } -}`, - 'T> ()', - String.raw`genericFunction () { - if (T isInstanceOf string) { - console.log(T) - } else { - // Do nothing - } -}` - ) - - /** - * const rawStr = "|Foo";| - * const anotherStr = "Bar" - */ - await assertClosingSymbolsHandler( - 'const rawStr = "', - String.raw` -const anotherStr = "Bar";`, - 'Foo";', - String.raw`const rawStr = "Foo"; -const anotherStr = "Bar";` - ) - - /** - * function shouldReturnAhtmlDiv( { name } : Props) { - * if (!name) { - * return undefined - * } - * - * return ( - *
- * { name } - *
- * |) - * } - */ - await assertClosingSymbolsHandler( - String.raw`function shouldReturnAhtmlDiv( { name } : Props) { - if (!name) { - return undefined - } - - return ( -
- { name } -
`, - String.raw`function shouldReturnAhtmlDiv( { name } : Props) { - if (!name) { - return undefined - } - - return ( -
- { name } -
- ) -}` - ) - }) - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/util/codeParsingUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/codeParsingUtil.test.ts deleted file mode 100644 index 2a2ad8bb34e..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/util/codeParsingUtil.test.ts +++ /dev/null @@ -1,327 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - PlatformLanguageId, - extractClasses, - extractFunctions, - isTestFile, - utgLanguageConfigs, -} from 'aws-core-vscode/codewhisperer' -import assert from 'assert' -import { createTestWorkspaceFolder, toTextDocument } from 'aws-core-vscode/test' - -describe('RegexValidationForPython', () => { - it('should extract all function names from a python file content', () => { - // TODO: Replace this variable based testing to read content from File. - // const filePath = vscode.Uri.file('./testData/samplePython.py').fsPath; - // const fileContent = fs.readFileSync('./testData/samplePython.py' , 'utf-8'); - // const regex = /function\s+(\w+)/g; - - const result = extractFunctions(pythonFileContent, utgLanguageConfigs['python'].functionExtractionPattern) - assert.strictEqual(result.length, 13) - assert.deepStrictEqual(result, [ - 'hello_world', - 'add_numbers', - 'multiply_numbers', - 'sum_numbers', - 'divide_numbers', - '__init__', - 'add', - 'multiply', - 'square', - 'from_sum', - '__init__', - 'triple', - 'main', - ]) - }) - - it('should extract all class names from a file content', () => { - const result = extractClasses(pythonFileContent, utgLanguageConfigs['python'].classExtractionPattern) - assert.deepStrictEqual(result, ['Calculator']) - }) -}) - -describe('RegexValidationForJava', () => { - it('should extract all function names from a java file content', () => { - // TODO: Replace this variable based testing to read content from File. - // const filePath = vscode.Uri.file('./testData/samplePython.py').fsPath; - // const fileContent = fs.readFileSync('./testData/samplePython.py' , 'utf-8'); - // const regex = /function\s+(\w+)/g; - - const result = extractFunctions(javaFileContent, utgLanguageConfigs['java'].functionExtractionPattern) - assert.strictEqual(result.length, 5) - assert.deepStrictEqual(result, ['sayHello', 'doSomething', 'square', 'manager', 'ABCFUNCTION']) - }) - - it('should extract all class names from a java file content', () => { - const result = extractClasses(javaFileContent, utgLanguageConfigs['java'].classExtractionPattern) - assert.deepStrictEqual(result, ['Test']) - }) -}) - -describe('isTestFile', () => { - let testWsFolder: string - beforeEach(async function () { - testWsFolder = (await createTestWorkspaceFolder()).uri.fsPath - }) - - it('validate by file path', async function () { - const langs = new Map([ - ['java', '.java'], - ['python', '.py'], - ['typescript', '.ts'], - ['javascript', '.js'], - ['typescriptreact', '.tsx'], - ['javascriptreact', '.jsx'], - ]) - const testFilePathsWithoutExt = [ - '/test/MyClass', - '/test/my_class', - '/tst/MyClass', - '/tst/my_class', - '/tests/MyClass', - '/tests/my_class', - ] - - const srcFilePathsWithoutExt = [ - '/src/MyClass', - 'MyClass', - 'foo/bar/MyClass', - 'foo/my_class', - 'my_class', - 'anyFolderOtherThanTest/foo/myClass', - ] - - for (const [languageId, ext] of langs) { - const testFilePaths = testFilePathsWithoutExt.map((it) => it + ext) - for (const testFilePath of testFilePaths) { - const actual = await isTestFile(testFilePath, { languageId: languageId }) - assert.strictEqual(actual, true) - } - - const srcFilePaths = srcFilePathsWithoutExt.map((it) => it + ext) - for (const srcFilePath of srcFilePaths) { - const actual = await isTestFile(srcFilePath, { languageId: languageId }) - assert.strictEqual(actual, false) - } - } - }) - - async function assertIsTestFile( - fileNames: string[], - config: { languageId: PlatformLanguageId }, - expected: boolean - ) { - for (const fileName of fileNames) { - const document = await toTextDocument('', fileName, testWsFolder) - const actual = await isTestFile(document.uri.fsPath, { languageId: config.languageId }) - assert.strictEqual(actual, expected) - } - } - - it('validate by file name', async function () { - const camelCaseSrc = ['Foo.java', 'Bar.java', 'Baz.java'] - await assertIsTestFile(camelCaseSrc, { languageId: 'java' }, false) - - const camelCaseTst = ['FooTest.java', 'BarTests.java'] - await assertIsTestFile(camelCaseTst, { languageId: 'java' }, true) - - const snakeCaseSrc = ['foo.py', 'bar.py'] - await assertIsTestFile(snakeCaseSrc, { languageId: 'python' }, false) - - const snakeCaseTst = ['test_foo.py', 'bar_test.py'] - await assertIsTestFile(snakeCaseTst, { languageId: 'python' }, true) - - const javascriptSrc = ['Foo.js', 'bar.js'] - await assertIsTestFile(javascriptSrc, { languageId: 'javascript' }, false) - - const javascriptTst = ['Foo.test.js', 'Bar.spec.js'] - await assertIsTestFile(javascriptTst, { languageId: 'javascript' }, true) - - const typescriptSrc = ['Foo.ts', 'bar.ts'] - await assertIsTestFile(typescriptSrc, { languageId: 'typescript' }, false) - - const typescriptTst = ['Foo.test.ts', 'Bar.spec.ts'] - await assertIsTestFile(typescriptTst, { languageId: 'typescript' }, true) - - const jsxSrc = ['Foo.jsx', 'Bar.jsx'] - await assertIsTestFile(jsxSrc, { languageId: 'javascriptreact' }, false) - - const jsxTst = ['Foo.test.jsx', 'Bar.spec.jsx'] - await assertIsTestFile(jsxTst, { languageId: 'javascriptreact' }, true) - }) - - it('should return true if the file name matches the test filename pattern - Java', async () => { - const filePaths = ['/path/to/MyClassTest.java', '/path/to/TestMyClass.java', '/path/to/MyClassTests.java'] - const language = 'java' - - for (const filePath of filePaths) { - const result = await isTestFile(filePath, { languageId: language }) - assert.strictEqual(result, true) - } - }) - - it('should return false if the file name does not match the test filename pattern - Java', async () => { - const filePaths = ['/path/to/MyClass.java', '/path/to/MyClass_test.java', '/path/to/test_MyClass.java'] - const language = 'java' - - for (const filePath of filePaths) { - const result = await isTestFile(filePath, { languageId: language }) - assert.strictEqual(result, false) - } - }) - - it('should return true if the file name does not match the test filename pattern - Python', async () => { - const filePaths = ['/path/to/util_test.py', '/path/to/test_util.py'] - const language = 'python' - - for (const filePath of filePaths) { - const result = await isTestFile(filePath, { languageId: language }) - assert.strictEqual(result, true) - } - }) - - it('should return false if the file name does not match the test filename pattern - Python', async () => { - const filePaths = ['/path/to/util.py', '/path/to/utilTest.java', '/path/to/Testutil.java'] - const language = 'python' - - for (const filePath of filePaths) { - const result = await isTestFile(filePath, { languageId: language }) - assert.strictEqual(result, false) - } - }) - - it('should return false if the language is not supported', async () => { - const filePath = '/path/to/MyClass.cpp' - const language = 'c++' - const result = await isTestFile(filePath, { languageId: language }) - assert.strictEqual(result, false) - }) -}) - -const pythonFileContent = ` -# Single-line import statements -import os -import numpy as np -from typing import List, Tuple - -# Multi-line import statements -from collections import ( - defaultdict, - Counter -) - -# Relative imports -from . import module1 -from ..subpackage import module2 - -# Wildcard imports -from mypackage import * -from mypackage.module import * - -# Aliased imports -import pandas as pd -from mypackage import module1 as m1, module2 as m2 - -def hello_world(): - print("Hello, world!") - -def add_numbers(x, y): - return x + y - -def multiply_numbers(x=1, y=1): - return x * y - -def sum_numbers(*args): - total = 0 - for num in args: - total += num - return total - -def divide_numbers(x, y=1, *args, **kwargs): - result = x / y - for arg in args: - result /= arg - for _, value in kwargs.items(): - result /= value - return result - -class Calculator: - def __init__(self, x, y): - self.x = x - self.y = y - - def add(self): - return self.x + self.y - - def multiply(self): - return self.x * self.y - - @staticmethod - def square(x): - return x ** 2 - - @classmethod - def from_sum(cls, x, y): - return cls(x+y, 0) - - class InnerClass: - def __init__(self, z): - self.z = z - - def triple(self): - return self.z * 3 - -def main(): - print(hello_world()) - print(add_numbers(3, 5)) - print(multiply_numbers(3, 5)) - print(sum_numbers(1, 2, 3, 4, 5)) - print(divide_numbers(10, 2, 5, 2, a=2, b=3)) - - calc = Calculator(3, 5) - print(calc.add()) - print(calc.multiply()) - print(Calculator.square(3)) - print(Calculator.from_sum(2, 3).add()) - - inner = Calculator.InnerClass(5) - print(inner.triple()) - -if __name__ == "__main__": - main() -` - -const javaFileContent = ` -@Annotation -public class Test { - Test() { - // Do something here - } - - //Additional commenting - public static void sayHello() { - System.out.println("Hello, World!"); - } - - private void doSomething(int x, int y) throws Exception { - int z = x + y; - System.out.println("The sum of " + x + " and " + y + " is " + z); - } - - protected static int square(int x) { - return x * x; - } - - private static void manager(int a, int b) { - return a+b; - } - - public int ABCFUNCTION( int ABC, int PQR) { - return ABC + PQR; - } -}` diff --git a/packages/amazonq/test/unit/codewhisperer/util/commonUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/commonUtil.test.ts deleted file mode 100644 index 5694b33365d..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/util/commonUtil.test.ts +++ /dev/null @@ -1,81 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import { - JsonConfigFileNamingConvention, - checkLeftContextKeywordsForJson, - getPrefixSuffixOverlap, -} from 'aws-core-vscode/codewhisperer' - -describe('commonUtil', function () { - describe('getPrefixSuffixOverlap', function () { - it('Should return correct overlap', async function () { - assert.strictEqual(getPrefixSuffixOverlap('32rasdgvdsg', 'sg462ydfgbs'), `sg`) - assert.strictEqual(getPrefixSuffixOverlap('32rasdgbreh', 'brehsega'), `breh`) - assert.strictEqual(getPrefixSuffixOverlap('42y24hsd', '42y24hsdzqq23'), `42y24hsd`) - assert.strictEqual(getPrefixSuffixOverlap('ge23yt1', 'ge23yt1'), `ge23yt1`) - assert.strictEqual(getPrefixSuffixOverlap('1sgdbsfbwsergsa', 'a1sgdbsfbwsergs'), `a`) - assert.strictEqual(getPrefixSuffixOverlap('xxa', 'xa'), `xa`) - }) - - it('Should return empty overlap for prefix suffix not matching cases', async function () { - assert.strictEqual(getPrefixSuffixOverlap('1sgdbsfbwsergsa', '1sgdbsfbwsergs'), ``) - assert.strictEqual(getPrefixSuffixOverlap('1sgdbsfbwsergsab', '1sgdbsfbwsergs'), ``) - assert.strictEqual(getPrefixSuffixOverlap('2135t12', 'v2135t12'), ``) - assert.strictEqual(getPrefixSuffixOverlap('2135t12', 'zv2135t12'), ``) - assert.strictEqual(getPrefixSuffixOverlap('xa', 'xxa'), ``) - }) - - it('Should return empty overlap for empty string input', async function () { - assert.strictEqual(getPrefixSuffixOverlap('ergwsghws', ''), ``) - assert.strictEqual(getPrefixSuffixOverlap('', 'asfegw4eh'), ``) - }) - }) - - describe('checkLeftContextKeywordsForJson', function () { - it('Should return true for valid left context keywords', async function () { - assert.strictEqual( - checkLeftContextKeywordsForJson('foo.json', 'Create an S3 Bucket named CodeWhisperer', 'json'), - true - ) - }) - it('Should return false for invalid left context keywords', async function () { - assert.strictEqual( - checkLeftContextKeywordsForJson( - 'foo.json', - 'Create an S3 Bucket named CodeWhisperer in Cloudformation', - 'json' - ), - false - ) - }) - - for (const jsonConfigFile of JsonConfigFileNamingConvention) { - it(`should evalute by filename ${jsonConfigFile}`, function () { - assert.strictEqual(checkLeftContextKeywordsForJson(jsonConfigFile, 'foo', 'json'), false) - - assert.strictEqual(checkLeftContextKeywordsForJson(jsonConfigFile.toUpperCase(), 'bar', 'json'), false) - - assert.strictEqual(checkLeftContextKeywordsForJson(jsonConfigFile.toUpperCase(), 'baz', 'json'), false) - }) - - const upperCaseFilename = jsonConfigFile.toUpperCase() - it(`should evalute by filename and case insensitive ${upperCaseFilename}`, function () { - assert.strictEqual(checkLeftContextKeywordsForJson(upperCaseFilename, 'foo', 'json'), false) - - assert.strictEqual( - checkLeftContextKeywordsForJson(upperCaseFilename.toUpperCase(), 'bar', 'json'), - false - ) - - assert.strictEqual( - checkLeftContextKeywordsForJson(upperCaseFilename.toUpperCase(), 'baz', 'json'), - false - ) - }) - } - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/util/crossFileContextUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/crossFileContextUtil.test.ts deleted file mode 100644 index 91e26e36111..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/util/crossFileContextUtil.test.ts +++ /dev/null @@ -1,454 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as FakeTimers from '@sinonjs/fake-timers' -import * as vscode from 'vscode' -import * as sinon from 'sinon' -import * as crossFile from 'aws-core-vscode/codewhisperer' -import { - aLongStringWithLineCount, - aStringWithLineCount, - createMockTextEditor, - installFakeClock, -} from 'aws-core-vscode/test' -import { FeatureConfigProvider, crossFileContextConfig } from 'aws-core-vscode/codewhisperer' -import { - assertTabCount, - closeAllEditors, - createTestWorkspaceFolder, - toTextEditor, - shuffleList, - toFile, -} from 'aws-core-vscode/test' -import { areEqual, normalize } from 'aws-core-vscode/shared' -import * as path from 'path' -import { LspController } from 'aws-core-vscode/amazonq' - -let tempFolder: string - -describe('crossFileContextUtil', function () { - const fakeCancellationToken: vscode.CancellationToken = { - isCancellationRequested: false, - onCancellationRequested: sinon.spy(), - } - - let mockEditor: vscode.TextEditor - let clock: FakeTimers.InstalledClock - - before(function () { - clock = installFakeClock() - }) - - after(function () { - clock.uninstall() - }) - - afterEach(function () { - sinon.restore() - }) - - describe('fetchSupplementalContextForSrc', function () { - beforeEach(async function () { - tempFolder = (await createTestWorkspaceFolder()).uri.fsPath - }) - - afterEach(async function () { - sinon.restore() - }) - - it.skip('for control group, should return opentabs context where there will be 3 chunks and each chunk should contains 50 lines', async function () { - sinon.stub(FeatureConfigProvider.instance, 'getProjectContextGroup').returns('control') - await toTextEditor(aStringWithLineCount(200), 'CrossFile.java', tempFolder, { preview: false }) - const myCurrentEditor = await toTextEditor('', 'TargetFile.java', tempFolder, { - preview: false, - }) - - await assertTabCount(2) - - const actual = await crossFile.fetchSupplementalContextForSrc(myCurrentEditor, fakeCancellationToken) - assert.ok(actual) - assert.strictEqual(actual.supplementalContextItems.length, 3) - assert.strictEqual(actual.supplementalContextItems[0].content.split('\n').length, 50) - assert.strictEqual(actual.supplementalContextItems[1].content.split('\n').length, 50) - assert.strictEqual(actual.supplementalContextItems[2].content.split('\n').length, 50) - }) - - it('for t1 group, should return repomap + opentabs context, should not exceed 20k total length', async function () { - await toTextEditor(aLongStringWithLineCount(200), 'CrossFile.java', tempFolder, { preview: false }) - const myCurrentEditor = await toTextEditor('', 'TargetFile.java', tempFolder, { - preview: false, - }) - - await assertTabCount(2) - - sinon.stub(FeatureConfigProvider.instance, 'getProjectContextGroup').returns('t1') - sinon - .stub(LspController.instance, 'queryInlineProjectContext') - .withArgs(sinon.match.any, sinon.match.any, 'codemap') - .resolves([ - { - content: 'foo'.repeat(3000), - score: 0, - filePath: 'q-inline', - }, - ]) - - const actual = await crossFile.fetchSupplementalContextForSrc(myCurrentEditor, fakeCancellationToken) - assert.ok(actual) - assert.strictEqual(actual.supplementalContextItems.length, 3) - assert.strictEqual(actual?.strategy, 'codemap') - assert.deepEqual(actual?.supplementalContextItems[0], { - content: 'foo'.repeat(3000), - score: 0, - filePath: 'q-inline', - }) - assert.strictEqual(actual.supplementalContextItems[1].content.split('\n').length, 50) - assert.strictEqual(actual.supplementalContextItems[2].content.split('\n').length, 50) - }) - - it.skip('for t2 group, should return global bm25 context and no repomap', async function () { - await toTextEditor(aStringWithLineCount(200), 'CrossFile.java', tempFolder, { preview: false }) - const myCurrentEditor = await toTextEditor('', 'TargetFile.java', tempFolder, { - preview: false, - }) - - await assertTabCount(2) - - sinon.stub(FeatureConfigProvider.instance, 'getProjectContextGroup').returns('t2') - sinon - .stub(LspController.instance, 'queryInlineProjectContext') - .withArgs(sinon.match.any, sinon.match.any, 'bm25') - .resolves([ - { - content: 'foo', - score: 5, - filePath: 'foo.java', - }, - { - content: 'bar', - score: 4, - filePath: 'bar.java', - }, - { - content: 'baz', - score: 3, - filePath: 'baz.java', - }, - { - content: 'qux', - score: 2, - filePath: 'qux.java', - }, - { - content: 'quux', - score: 1, - filePath: 'quux.java', - }, - ]) - - const actual = await crossFile.fetchSupplementalContextForSrc(myCurrentEditor, fakeCancellationToken) - assert.ok(actual) - assert.strictEqual(actual.supplementalContextItems.length, 5) - assert.strictEqual(actual?.strategy, 'bm25') - - assert.deepEqual(actual?.supplementalContextItems[0], { - content: 'foo', - score: 5, - filePath: 'foo.java', - }) - - assert.deepEqual(actual?.supplementalContextItems[1], { - content: 'bar', - score: 4, - filePath: 'bar.java', - }) - assert.deepEqual(actual?.supplementalContextItems[2], { - content: 'baz', - score: 3, - filePath: 'baz.java', - }) - - assert.deepEqual(actual?.supplementalContextItems[3], { - content: 'qux', - score: 2, - filePath: 'qux.java', - }) - - assert.deepEqual(actual?.supplementalContextItems[4], { - content: 'quux', - score: 1, - filePath: 'quux.java', - }) - }) - }) - - describe('non supported language should return undefined', function () { - it('c++', async function () { - mockEditor = createMockTextEditor('content', 'fileName', 'cpp') - const actual = await crossFile.fetchSupplementalContextForSrc(mockEditor, fakeCancellationToken) - assert.strictEqual(actual, undefined) - }) - - it('ruby', async function () { - mockEditor = createMockTextEditor('content', 'fileName', 'ruby') - - const actual = await crossFile.fetchSupplementalContextForSrc(mockEditor, fakeCancellationToken) - - assert.strictEqual(actual, undefined) - }) - }) - - describe('getCrossFileCandidate', function () { - before(async function () { - this.timeout(60000) - }) - - beforeEach(async function () { - tempFolder = (await createTestWorkspaceFolder()).uri.fsPath - }) - - afterEach(async function () { - await closeAllEditors() - }) - - it('should return opened files, exclude test files and sorted ascendingly by file distance', async function () { - const targetFile = path.join('src', 'service', 'microService', 'CodeWhispererFileContextProvider.java') - const fileWithDistance3 = path.join('src', 'service', 'CodewhispererRecommendationService.java') - const fileWithDistance5 = path.join('src', 'util', 'CodeWhispererConstants.java') - const fileWithDistance6 = path.join('src', 'ui', 'popup', 'CodeWhispererPopupManager.java') - const fileWithDistance7 = path.join('src', 'ui', 'popup', 'components', 'CodeWhispererPopup.java') - const fileWithDistance8 = path.join( - 'src', - 'ui', - 'popup', - 'components', - 'actions', - 'AcceptRecommendationAction.java' - ) - const testFile1 = path.join('test', 'service', 'CodeWhispererFileContextProviderTest.java') - const testFile2 = path.join('test', 'ui', 'CodeWhispererPopupManagerTest.java') - - const expectedFilePaths = [ - fileWithDistance3, - fileWithDistance5, - fileWithDistance6, - fileWithDistance7, - fileWithDistance8, - ] - - const shuffledFilePaths = shuffleList(expectedFilePaths) - - for (const filePath of shuffledFilePaths) { - await toTextEditor('', filePath, tempFolder, { preview: false }) - } - - await toTextEditor('', testFile1, tempFolder, { preview: false }) - await toTextEditor('', testFile2, tempFolder, { preview: false }) - const editor = await toTextEditor('', targetFile, tempFolder, { preview: false }) - - await assertTabCount(shuffledFilePaths.length + 3) - - const actual = await crossFile.getCrossFileCandidates(editor) - - assert.ok(actual.length === 5) - for (const [index, actualFile] of actual.entries()) { - const expectedFile = path.join(tempFolder, expectedFilePaths[index]) - assert.strictEqual(normalize(expectedFile), normalize(actualFile)) - assert.ok(areEqual(tempFolder, actualFile, expectedFile)) - } - }) - }) - - describe.skip('partial support - control group', function () { - const fileExtLists: string[] = [] - - before(async function () { - this.timeout(60000) - }) - - beforeEach(async function () { - tempFolder = (await createTestWorkspaceFolder()).uri.fsPath - }) - - afterEach(async function () { - await closeAllEditors() - }) - - for (const fileExt of fileExtLists) { - it('should be empty if userGroup is control', async function () { - const editor = await toTextEditor('content-1', `file-1.${fileExt}`, tempFolder) - await toTextEditor('content-2', `file-2.${fileExt}`, tempFolder, { preview: false }) - await toTextEditor('content-3', `file-3.${fileExt}`, tempFolder, { preview: false }) - await toTextEditor('content-4', `file-4.${fileExt}`, tempFolder, { preview: false }) - - await assertTabCount(4) - - const actual = await crossFile.fetchSupplementalContextForSrc(editor, fakeCancellationToken) - - assert.ok(actual && actual.supplementalContextItems.length === 0) - }) - } - }) - - describe.skip('partial support - crossfile group', function () { - const fileExtLists: string[] = [] - - before(async function () { - this.timeout(60000) - }) - - beforeEach(async function () { - tempFolder = (await createTestWorkspaceFolder()).uri.fsPath - }) - - afterEach(async function () { - await closeAllEditors() - }) - - for (const fileExt of fileExtLists) { - it('should be non empty if usergroup is Crossfile', async function () { - const editor = await toTextEditor('content-1', `file-1.${fileExt}`, tempFolder) - await toTextEditor('content-2', `file-2.${fileExt}`, tempFolder, { preview: false }) - await toTextEditor('content-3', `file-3.${fileExt}`, tempFolder, { preview: false }) - await toTextEditor('content-4', `file-4.${fileExt}`, tempFolder, { preview: false }) - - await assertTabCount(4) - - const actual = await crossFile.fetchSupplementalContextForSrc(editor, fakeCancellationToken) - - assert.ok(actual && actual.supplementalContextItems.length !== 0) - }) - } - }) - - describe('full support', function () { - const fileExtLists = ['java', 'js', 'ts', 'py', 'tsx', 'jsx'] - - before(async function () { - this.timeout(60000) - }) - - beforeEach(async function () { - tempFolder = (await createTestWorkspaceFolder()).uri.fsPath - }) - - afterEach(async function () { - sinon.restore() - await closeAllEditors() - }) - - for (const fileExt of fileExtLists) { - it(`supplemental context for file ${fileExt} should be non empty`, async function () { - sinon.stub(FeatureConfigProvider.instance, 'getProjectContextGroup').returns('control') - sinon - .stub(LspController.instance, 'queryInlineProjectContext') - .withArgs(sinon.match.any, sinon.match.any, 'codemap') - .resolves([ - { - content: 'foo', - score: 0, - filePath: 'q-inline', - }, - ]) - const editor = await toTextEditor('content-1', `file-1.${fileExt}`, tempFolder) - await toTextEditor('content-2', `file-2.${fileExt}`, tempFolder, { preview: false }) - await toTextEditor('content-3', `file-3.${fileExt}`, tempFolder, { preview: false }) - await toTextEditor('content-4', `file-4.${fileExt}`, tempFolder, { preview: false }) - - await assertTabCount(4) - - const actual = await crossFile.fetchSupplementalContextForSrc(editor, fakeCancellationToken) - - assert.ok(actual && actual.supplementalContextItems.length !== 0) - }) - } - }) - - describe('splitFileToChunks', function () { - beforeEach(async function () { - tempFolder = (await createTestWorkspaceFolder()).uri.fsPath - }) - - it('should split file to a chunk of 2 lines', async function () { - const filePath = path.join(tempFolder, 'file.txt') - await toFile('line_1\nline_2\nline_3\nline_4\nline_5\nline_6\nline_7', filePath) - - const chunks = await crossFile.splitFileToChunks(filePath, 2) - - assert.strictEqual(chunks.length, 4) - assert.strictEqual(chunks[0].content, 'line_1\nline_2') - assert.strictEqual(chunks[1].content, 'line_3\nline_4') - assert.strictEqual(chunks[2].content, 'line_5\nline_6') - assert.strictEqual(chunks[3].content, 'line_7') - }) - - it('should split file to a chunk of 5 lines', async function () { - const filePath = path.join(tempFolder, 'file.txt') - await toFile('line_1\nline_2\nline_3\nline_4\nline_5\nline_6\nline_7', filePath) - - const chunks = await crossFile.splitFileToChunks(filePath, 5) - - assert.strictEqual(chunks.length, 2) - assert.strictEqual(chunks[0].content, 'line_1\nline_2\nline_3\nline_4\nline_5') - assert.strictEqual(chunks[1].content, 'line_6\nline_7') - }) - - it('codewhisperer crossfile config should use 50 lines', async function () { - const filePath = path.join(tempFolder, 'file.txt') - await toFile(aStringWithLineCount(210), filePath) - - const chunks = await crossFile.splitFileToChunks(filePath, crossFileContextConfig.numberOfLinesEachChunk) - - // (210 / 50) + 1 - assert.strictEqual(chunks.length, 5) - // line0 -> line49 - assert.strictEqual(chunks[0].content, aStringWithLineCount(50, 0)) - // line50 -> line99 - assert.strictEqual(chunks[1].content, aStringWithLineCount(50, 50)) - // line100 -> line149 - assert.strictEqual(chunks[2].content, aStringWithLineCount(50, 100)) - // line150 -> line199 - assert.strictEqual(chunks[3].content, aStringWithLineCount(50, 150)) - // line 200 -> line209 - assert.strictEqual(chunks[4].content, aStringWithLineCount(10, 200)) - }) - - it('linkChunks should add another chunk which will link to the first chunk and chunk.nextContent should reflect correct value', async function () { - const filePath = path.join(tempFolder, 'file.txt') - await toFile(aStringWithLineCount(210), filePath) - - const chunks = await crossFile.splitFileToChunks(filePath, crossFileContextConfig.numberOfLinesEachChunk) - const linkedChunks = crossFile.linkChunks(chunks) - - // 210 / 50 + 2 - assert.strictEqual(linkedChunks.length, 6) - - // 0th - assert.strictEqual(linkedChunks[0].content, aStringWithLineCount(3, 0)) - assert.strictEqual(linkedChunks[0].nextContent, aStringWithLineCount(50, 0)) - - // 1st - assert.strictEqual(linkedChunks[1].content, aStringWithLineCount(50, 0)) - assert.strictEqual(linkedChunks[1].nextContent, aStringWithLineCount(50, 50)) - - // 2nd - assert.strictEqual(linkedChunks[2].content, aStringWithLineCount(50, 50)) - assert.strictEqual(linkedChunks[2].nextContent, aStringWithLineCount(50, 100)) - - // 3rd - assert.strictEqual(linkedChunks[3].content, aStringWithLineCount(50, 100)) - assert.strictEqual(linkedChunks[3].nextContent, aStringWithLineCount(50, 150)) - - // 4th - assert.strictEqual(linkedChunks[4].content, aStringWithLineCount(50, 150)) - assert.strictEqual(linkedChunks[4].nextContent, aStringWithLineCount(10, 200)) - - // 5th - assert.strictEqual(linkedChunks[5].content, aStringWithLineCount(10, 200)) - assert.strictEqual(linkedChunks[5].nextContent, aStringWithLineCount(10, 200)) - }) - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts deleted file mode 100644 index 3875dbbd0f2..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts +++ /dev/null @@ -1,392 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -import assert from 'assert' -import * as codewhispererClient from 'aws-core-vscode/codewhisperer' -import * as EditorContext from 'aws-core-vscode/codewhisperer' -import { - createMockDocument, - createMockTextEditor, - createMockClientRequest, - resetCodeWhispererGlobalVariables, - toTextEditor, - createTestWorkspaceFolder, - closeAllEditors, -} from 'aws-core-vscode/test' -import { globals } from 'aws-core-vscode/shared' -import { GenerateCompletionsRequest } from 'aws-core-vscode/codewhisperer' -import * as vscode from 'vscode' - -export function createNotebookCell( - document: vscode.TextDocument = createMockDocument('def example():\n return "test"'), - kind: vscode.NotebookCellKind = vscode.NotebookCellKind.Code, - notebook: vscode.NotebookDocument = {} as any, - index: number = 0, - outputs: vscode.NotebookCellOutput[] = [], - metadata: { readonly [key: string]: any } = {}, - executionSummary?: vscode.NotebookCellExecutionSummary -): vscode.NotebookCell { - return { - document, - kind, - notebook, - index, - outputs, - metadata, - executionSummary, - } -} - -describe('editorContext', function () { - let telemetryEnabledDefault: boolean - let tempFolder: string - - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - telemetryEnabledDefault = globals.telemetry.telemetryEnabled - }) - - afterEach(async function () { - await globals.telemetry.setTelemetryEnabled(telemetryEnabledDefault) - }) - - describe('extractContextForCodeWhisperer', function () { - it('Should return expected context', function () { - const editor = createMockTextEditor('import math\ndef two_sum(nums, target):\n', 'test.py', 'python', 1, 17) - const actual = EditorContext.extractContextForCodeWhisperer(editor) - const expected: codewhispererClient.FileContext = { - fileUri: 'file:///test.py', - filename: 'test.py', - programmingLanguage: { - languageName: 'python', - }, - leftFileContent: 'import math\ndef two_sum(nums,', - rightFileContent: ' target):\n', - } - assert.deepStrictEqual(actual, expected) - }) - - it('Should return expected context within max char limit', function () { - const editor = createMockTextEditor( - 'import math\ndef ' + 'a'.repeat(10340) + 'two_sum(nums, target):\n', - 'test.py', - 'python', - 1, - 17 - ) - const actual = EditorContext.extractContextForCodeWhisperer(editor) - const expected: codewhispererClient.FileContext = { - fileUri: 'file:///test.py', - filename: 'test.py', - programmingLanguage: { - languageName: 'python', - }, - leftFileContent: 'import math\ndef aaaaaaaaaaaaa', - rightFileContent: 'a'.repeat(10240), - } - assert.deepStrictEqual(actual, expected) - }) - - it('in a notebook, includes context from other cells', async function () { - const cells: vscode.NotebookCellData[] = [ - new vscode.NotebookCellData(vscode.NotebookCellKind.Markup, 'Previous cell', 'python'), - new vscode.NotebookCellData( - vscode.NotebookCellKind.Code, - 'import numpy as np\nimport pandas as pd\n\ndef analyze_data(df):\n # Current cell with cursor here', - 'python' - ), - new vscode.NotebookCellData( - vscode.NotebookCellKind.Code, - '# Process the data\nresult = analyze_data(df)\nprint(result)', - 'python' - ), - ] - - const document = await vscode.workspace.openNotebookDocument( - 'jupyter-notebook', - new vscode.NotebookData(cells) - ) - const editor: any = { - document: document.cellAt(1).document, - selection: { active: new vscode.Position(4, 13) }, - } - - const actual = EditorContext.extractContextForCodeWhisperer(editor) - const expected: codewhispererClient.FileContext = { - fileUri: editor.document.uri.toString(), - filename: 'Untitled-1.py', - programmingLanguage: { - languageName: 'python', - }, - leftFileContent: - '# Previous cell\nimport numpy as np\nimport pandas as pd\n\ndef analyze_data(df):\n # Current', - rightFileContent: - ' cell with cursor here\n# Process the data\nresult = analyze_data(df)\nprint(result)\n', - } - assert.deepStrictEqual(actual, expected) - }) - }) - - describe('getFileName', function () { - it('Should return expected filename given a document reading test.py', function () { - const editor = createMockTextEditor('', 'test.py', 'python', 1, 17) - const actual = EditorContext.getFileName(editor) - const expected = 'test.py' - assert.strictEqual(actual, expected) - }) - - it('Should return expected filename for a long filename', async function () { - const editor = createMockTextEditor('', 'a'.repeat(1500), 'python', 1, 17) - const actual = EditorContext.getFileName(editor) - const expected = 'a'.repeat(1024) - assert.strictEqual(actual, expected) - }) - }) - - describe('getFileRelativePath', function () { - this.beforeEach(async function () { - tempFolder = (await createTestWorkspaceFolder()).uri.fsPath - }) - - it('Should return a new filename with correct extension given a .ipynb file', function () { - const languageToExtension = new Map([ - ['python', 'py'], - ['rust', 'rs'], - ['javascript', 'js'], - ['typescript', 'ts'], - ['c', 'c'], - ]) - - for (const [language, extension] of languageToExtension.entries()) { - const editor = createMockTextEditor('', 'test.ipynb', language, 1, 17) - const actual = EditorContext.getFileRelativePath(editor) - const expected = 'test.' + extension - assert.strictEqual(actual, expected) - } - }) - - it('Should return relative path', async function () { - const editor = await toTextEditor('tttt', 'test.py', tempFolder) - const actual = EditorContext.getFileRelativePath(editor) - const expected = 'test.py' - assert.strictEqual(actual, expected) - }) - - afterEach(async function () { - await closeAllEditors() - }) - }) - - describe('getNotebookCellContext', function () { - it('Should return cell text for python code cells when language is python', function () { - const mockCodeCell = createNotebookCell(createMockDocument('def example():\n return "test"')) - const result = EditorContext.getNotebookCellContext(mockCodeCell, 'python') - assert.strictEqual(result, 'def example():\n return "test"') - }) - - it('Should return java comments for python code cells when language is java', function () { - const mockCodeCell = createNotebookCell(createMockDocument('def example():\n return "test"')) - const result = EditorContext.getNotebookCellContext(mockCodeCell, 'java') - assert.strictEqual(result, '// def example():\n// return "test"') - }) - - it('Should return python comments for java code cells when language is python', function () { - const mockCodeCell = createNotebookCell(createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java')) - const result = EditorContext.getNotebookCellContext(mockCodeCell, 'python') - assert.strictEqual(result, '# println(1 + 1);') - }) - - it('Should add python comment prefixes for markdown cells when language is python', function () { - const mockMarkdownCell = createNotebookCell( - createMockDocument('# Heading\nThis is a markdown cell'), - vscode.NotebookCellKind.Markup - ) - const result = EditorContext.getNotebookCellContext(mockMarkdownCell, 'python') - assert.strictEqual(result, '# # Heading\n# This is a markdown cell') - }) - - it('Should add java comment prefixes for markdown cells when language is java', function () { - const mockMarkdownCell = createNotebookCell( - createMockDocument('# Heading\nThis is a markdown cell'), - vscode.NotebookCellKind.Markup - ) - const result = EditorContext.getNotebookCellContext(mockMarkdownCell, 'java') - assert.strictEqual(result, '// # Heading\n// This is a markdown cell') - }) - }) - - describe('getNotebookCellsSliceContext', function () { - it('Should extract content from cells in reverse order up to maxLength from prefix cells', function () { - const mockCells = [ - createNotebookCell(createMockDocument('First cell content')), - createNotebookCell(createMockDocument('Second cell content')), - createNotebookCell(createMockDocument('Third cell content')), - ] - - const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', false) - assert.strictEqual(result, 'First cell content\nSecond cell content\nThird cell content\n') - }) - - it('Should extract content from cells in reverse order up to maxLength from suffix cells', function () { - const mockCells = [ - createNotebookCell(createMockDocument('First cell content')), - createNotebookCell(createMockDocument('Second cell content')), - createNotebookCell(createMockDocument('Third cell content')), - ] - - const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', true) - assert.strictEqual(result, 'First cell content\nSecond cell content\nThird cell content\n') - }) - - it('Should respect maxLength parameter from prefix cells', function () { - const mockCells = [ - createNotebookCell(createMockDocument('First')), - createNotebookCell(createMockDocument('Second')), - createNotebookCell(createMockDocument('Third')), - createNotebookCell(createMockDocument('Fourth')), - ] - // Should only include part of second cell and the last two cells - const result = EditorContext.getNotebookCellsSliceContext(mockCells, 15, 'python', false) - assert.strictEqual(result, 'd\nThird\nFourth\n') - }) - - it('Should respect maxLength parameter from suffix cells', function () { - const mockCells = [ - createNotebookCell(createMockDocument('First')), - createNotebookCell(createMockDocument('Second')), - createNotebookCell(createMockDocument('Third')), - createNotebookCell(createMockDocument('Fourth')), - ] - - // Should only include first cell and part of second cell - const result = EditorContext.getNotebookCellsSliceContext(mockCells, 15, 'python', true) - assert.strictEqual(result, 'First\nSecond\nTh') - }) - - it('Should handle empty cells array from prefix cells', function () { - const result = EditorContext.getNotebookCellsSliceContext([], 100, 'python', false) - assert.strictEqual(result, '') - }) - - it('Should handle empty cells array from suffix cells', function () { - const result = EditorContext.getNotebookCellsSliceContext([], 100, 'python', true) - assert.strictEqual(result, '') - }) - - it('Should add python comments to markdown prefix cells', function () { - const mockCells = [ - createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), - createNotebookCell(createMockDocument('def example():\n return "test"')), - ] - const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', false) - assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n return "test"\n') - }) - - it('Should add python comments to markdown suffix cells', function () { - const mockCells = [ - createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), - createNotebookCell(createMockDocument('def example():\n return "test"')), - ] - - const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', true) - assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n return "test"\n') - }) - - it('Should add java comments to markdown and python prefix cells when language is java', function () { - const mockCells = [ - createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), - createNotebookCell(createMockDocument('def example():\n return "test"')), - ] - const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'java', false) - assert.strictEqual(result, '// # Heading\n// This is markdown\n// def example():\n// return "test"\n') - }) - - it('Should add java comments to markdown and python suffix cells when language is java', function () { - const mockCells = [ - createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), - createNotebookCell(createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java')), - ] - - const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'java', true) - assert.strictEqual(result, '// # Heading\n// This is markdown\nprintln(1 + 1);\n') - }) - - it('Should handle code prefix cells with different languages', function () { - const mockCells = [ - createNotebookCell( - createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java'), - vscode.NotebookCellKind.Code - ), - createNotebookCell(createMockDocument('def example():\n return "test"')), - ] - const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', false) - assert.strictEqual(result, '# println(1 + 1);\ndef example():\n return "test"\n') - }) - - it('Should handle code suffix cells with different languages', function () { - const mockCells = [ - createNotebookCell( - createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java'), - vscode.NotebookCellKind.Code - ), - createNotebookCell(createMockDocument('def example():\n return "test"')), - ] - const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', true) - assert.strictEqual(result, '# println(1 + 1);\ndef example():\n return "test"\n') - }) - }) - - describe('validateRequest', function () { - it('Should return false if request filename.length is invalid', function () { - const req = createMockClientRequest() - req.fileContext.filename = '' - assert.ok(!EditorContext.validateRequest(req)) - }) - - it('Should return false if request programming language is invalid', function () { - const req = createMockClientRequest() - req.fileContext.programmingLanguage.languageName = '' - assert.ok(!EditorContext.validateRequest(req)) - req.fileContext.programmingLanguage.languageName = 'a'.repeat(200) - assert.ok(!EditorContext.validateRequest(req)) - }) - - it('Should return false if request left or right context exceeds max length', function () { - const req = createMockClientRequest() - req.fileContext.leftFileContent = 'a'.repeat(256000) - assert.ok(!EditorContext.validateRequest(req)) - req.fileContext.leftFileContent = 'a' - req.fileContext.rightFileContent = 'a'.repeat(256000) - assert.ok(!EditorContext.validateRequest(req)) - }) - - it('Should return true if above conditions are not met', function () { - const req = createMockClientRequest() - assert.ok(EditorContext.validateRequest(req)) - }) - }) - - describe('getLeftContext', function () { - it('Should return expected left context', function () { - const editor = createMockTextEditor('import math\ndef two_sum(nums, target):\n', 'test.py', 'python', 1, 17) - const actual = EditorContext.getLeftContext(editor, 1) - const expected = '...wo_sum(nums, target)' - assert.strictEqual(actual, expected) - }) - }) - - describe('buildListRecommendationRequest', function () { - it('Should return expected fields for optOut, nextToken and reference config', async function () { - const nextToken = 'testToken' - const optOutPreference = false - await globals.telemetry.setTelemetryEnabled(false) - const editor = createMockTextEditor('import math\ndef two_sum(nums, target):\n', 'test.py', 'python', 1, 17) - const actual = await EditorContext.buildListRecommendationRequest(editor, nextToken, optOutPreference) - - assert.strictEqual(actual.request.nextToken, nextToken) - assert.strictEqual((actual.request as GenerateCompletionsRequest).optOutPreference, 'OPTOUT') - assert.strictEqual(actual.request.referenceTrackerConfiguration?.recommendationsWithReferences, 'BLOCK') - }) - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/util/globalStateUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/globalStateUtil.test.ts deleted file mode 100644 index 24062a81b7c..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/util/globalStateUtil.test.ts +++ /dev/null @@ -1,42 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as sinon from 'sinon' -import { resetCodeWhispererGlobalVariables } from 'aws-core-vscode/test' -import { getLogger } from 'aws-core-vscode/shared' -import { resetIntelliSenseState, vsCodeState } from 'aws-core-vscode/codewhisperer' - -describe('globalStateUtil', function () { - let loggerSpy: sinon.SinonSpy - - beforeEach(async function () { - await resetCodeWhispererGlobalVariables() - vsCodeState.isIntelliSenseActive = true - loggerSpy = sinon.spy(getLogger(), 'info') - }) - - this.afterEach(function () { - sinon.restore() - }) - - it('Should skip when CodeWhisperer is turned off', async function () { - const isManualTriggerEnabled = false - const isAutomatedTriggerEnabled = false - resetIntelliSenseState(isManualTriggerEnabled, isAutomatedTriggerEnabled, true) - assert.ok(!loggerSpy.called) - }) - - it('Should skip when invocationContext is not active', async function () { - vsCodeState.isIntelliSenseActive = false - resetIntelliSenseState(false, false, true) - assert.ok(!loggerSpy.called) - }) - - it('Should skip when no valid recommendations', async function () { - resetIntelliSenseState(true, true, false) - assert.ok(!loggerSpy.called) - }) -}) diff --git a/packages/amazonq/test/unit/codewhisperer/util/showSsoPrompt.test.ts b/packages/amazonq/test/unit/codewhisperer/util/showSsoPrompt.test.ts index 7d004e8ede5..1d67db60efc 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/showSsoPrompt.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/showSsoPrompt.test.ts @@ -7,12 +7,20 @@ import * as vscode from 'vscode' import assert from 'assert' import * as sinon from 'sinon' import { resetCodeWhispererGlobalVariables } from 'aws-core-vscode/test' -import { assertTelemetryCurried, getTestWindow, getTestLogger } from 'aws-core-vscode/test' +import { assertTelemetryCurried, getTestWindow } from 'aws-core-vscode/test' import { AuthUtil, awsIdSignIn, showCodeWhispererConnectionPrompt } from 'aws-core-vscode/codewhisperer' +import { SsoAccessTokenProvider, constants } from 'aws-core-vscode/auth' describe('showConnectionPrompt', function () { + let isBuilderIdConnection: sinon.SinonStub + beforeEach(async function () { await resetCodeWhispererGlobalVariables() + isBuilderIdConnection = sinon.stub(AuthUtil.instance, 'isBuilderIdConnection') + isBuilderIdConnection.resolves() + + // Stub useDeviceFlow so we always use DeviceFlow for auth + sinon.stub(SsoAccessTokenProvider, 'useDeviceFlow').returns(true) }) afterEach(function () { @@ -20,7 +28,7 @@ describe('showConnectionPrompt', function () { }) it('can select connect to AwsBuilderId', async function () { - const authUtilSpy = sinon.stub(AuthUtil.instance, 'connectToAwsBuilderId') + sinon.stub(AuthUtil.instance, 'login').resolves() getTestWindow().onDidShowQuickPick(async (picker) => { await picker.untilReady() @@ -29,18 +37,18 @@ describe('showConnectionPrompt', function () { await showCodeWhispererConnectionPrompt() - assert.ok(authUtilSpy.called) const assertTelemetry = assertTelemetryCurried('ui_click') assertTelemetry({ elementId: 'connection_optionBuilderID' }) + assert.ok(isBuilderIdConnection) }) - it('connectToAwsBuilderId logs that AWS ID sign in was selected', async function () { - sinon.stub(AuthUtil.instance, 'connectToAwsBuilderId').resolves() + it('connectToAwsBuilderId calls AuthUtil login with builderIdStartUrl', async function () { sinon.stub(vscode.commands, 'executeCommand') + const loginStub = sinon.stub(AuthUtil.instance, 'login').resolves() await awsIdSignIn() - const loggedEntries = getTestLogger().getLoggedEntries() - assert.ok(loggedEntries.find((entry) => entry === 'selected AWS ID sign in')) + assert.strictEqual(loginStub.called, true) + assert.strictEqual(loginStub.firstCall.args[0], constants.builderIdStartUrl) }) }) diff --git a/packages/amazonq/test/unit/codewhisperer/util/supplemetalContextUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/supplemetalContextUtil.test.ts deleted file mode 100644 index a42b0aa6158..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/util/supplemetalContextUtil.test.ts +++ /dev/null @@ -1,265 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as FakeTimers from '@sinonjs/fake-timers' -import * as vscode from 'vscode' -import * as sinon from 'sinon' -import * as os from 'os' -import * as crossFile from 'aws-core-vscode/codewhisperer' -import { TestFolder, assertTabCount, installFakeClock } from 'aws-core-vscode/test' -import { CodeWhispererSupplementalContext, FeatureConfigProvider } from 'aws-core-vscode/codewhisperer' -import { toTextEditor } from 'aws-core-vscode/test' -import { LspController } from 'aws-core-vscode/amazonq' - -const newLine = os.EOL - -describe('supplementalContextUtil', function () { - let testFolder: TestFolder - let clock: FakeTimers.InstalledClock - - const fakeCancellationToken: vscode.CancellationToken = { - isCancellationRequested: false, - onCancellationRequested: sinon.spy(), - } - - before(function () { - clock = installFakeClock() - }) - - after(function () { - clock.uninstall() - }) - - beforeEach(async function () { - testFolder = await TestFolder.create() - sinon.stub(FeatureConfigProvider.instance, 'getProjectContextGroup').returns('control') - }) - - afterEach(function () { - sinon.restore() - }) - - describe('fetchSupplementalContext', function () { - describe('openTabsContext', function () { - it('opentabContext should include chunks if non empty', async function () { - sinon - .stub(LspController.instance, 'queryInlineProjectContext') - .withArgs(sinon.match.any, sinon.match.any, 'codemap') - .resolves([ - { - content: 'foo', - score: 0, - filePath: 'q-inline', - }, - ]) - await toTextEditor('class Foo', 'Foo.java', testFolder.path, { preview: false }) - await toTextEditor('class Bar', 'Bar.java', testFolder.path, { preview: false }) - await toTextEditor('class Baz', 'Baz.java', testFolder.path, { preview: false }) - - const editor = await toTextEditor('public class Foo {}', 'Query.java', testFolder.path, { - preview: false, - }) - - await assertTabCount(4) - - const actual = await crossFile.fetchSupplementalContext(editor, fakeCancellationToken) - assert.ok(actual?.supplementalContextItems.length === 4) - }) - - it('opentabsContext should filter out empty chunks', async function () { - // open 3 files as supplemental context candidate files but none of them have contents - await toTextEditor('', 'Foo.java', testFolder.path, { preview: false }) - await toTextEditor('', 'Bar.java', testFolder.path, { preview: false }) - await toTextEditor('', 'Baz.java', testFolder.path, { preview: false }) - - const editor = await toTextEditor('public class Foo {}', 'Query.java', testFolder.path, { - preview: false, - }) - - await assertTabCount(4) - - const actual = await crossFile.fetchSupplementalContext(editor, fakeCancellationToken) - assert.ok(actual?.supplementalContextItems.length === 0) - }) - }) - }) - - describe('truncation', function () { - it('truncate context should do nothing if everything fits in constraint', function () { - const chunkA: crossFile.CodeWhispererSupplementalContextItem = { - content: 'a', - filePath: 'a.java', - score: 0, - } - const chunkB: crossFile.CodeWhispererSupplementalContextItem = { - content: 'b', - filePath: 'b.java', - score: 1, - } - const chunks = [chunkA, chunkB] - - const supplementalContext: CodeWhispererSupplementalContext = { - isUtg: false, - isProcessTimeout: false, - supplementalContextItems: chunks, - contentsLength: 25000, - latency: 0, - strategy: 'codemap', - } - - const actual = crossFile.truncateSuppelementalContext(supplementalContext) - assert.strictEqual(actual.supplementalContextItems.length, 2) - assert.strictEqual(actual.supplementalContextItems[0].content, 'a') - assert.strictEqual(actual.supplementalContextItems[1].content, 'b') - }) - - it('truncateLineByLine should drop the last line if max length is greater than threshold', function () { - const input = - repeatString('a', 11) + - newLine + - repeatString('b', 11) + - newLine + - repeatString('c', 11) + - newLine + - repeatString('d', 11) + - newLine + - repeatString('e', 11) - - assert.ok(input.length > 50) - const actual = crossFile.truncateLineByLine(input, 50) - assert.ok(actual.length <= 50) - - const input2 = repeatString(`b${newLine}`, 10) - const actual2 = crossFile.truncateLineByLine(input2, 8) - assert.ok(actual2.length <= 8) - }) - - it('truncation context should make context length per item lte 10240 cap', function () { - const chunkA: crossFile.CodeWhispererSupplementalContextItem = { - content: repeatString(`a${newLine}`, 4000), - filePath: 'a.java', - score: 0, - } - const chunkB: crossFile.CodeWhispererSupplementalContextItem = { - content: repeatString(`b${newLine}`, 6000), - filePath: 'b.java', - score: 1, - } - const chunkC: crossFile.CodeWhispererSupplementalContextItem = { - content: repeatString(`c${newLine}`, 1000), - filePath: 'c.java', - score: 2, - } - const chunkD: crossFile.CodeWhispererSupplementalContextItem = { - content: repeatString(`d${newLine}`, 1500), - filePath: 'd.java', - score: 3, - } - - assert.ok( - chunkA.content.length + chunkB.content.length + chunkC.content.length + chunkD.content.length > 20480 - ) - - const supplementalContext: CodeWhispererSupplementalContext = { - isUtg: false, - isProcessTimeout: false, - supplementalContextItems: [chunkA, chunkB, chunkC, chunkD], - contentsLength: 25000, - latency: 0, - strategy: 'codemap', - } - - const actual = crossFile.truncateSuppelementalContext(supplementalContext) - assert.strictEqual(actual.supplementalContextItems.length, 3) - assert.ok(actual.contentsLength <= 20480) - assert.strictEqual(actual.strategy, 'codemap') - }) - - it('truncate context should make context items lte 5', function () { - const chunkA: crossFile.CodeWhispererSupplementalContextItem = { - content: 'a', - filePath: 'a.java', - score: 0, - } - const chunkB: crossFile.CodeWhispererSupplementalContextItem = { - content: 'b', - filePath: 'b.java', - score: 1, - } - const chunkC: crossFile.CodeWhispererSupplementalContextItem = { - content: 'c', - filePath: 'c.java', - score: 2, - } - const chunkD: crossFile.CodeWhispererSupplementalContextItem = { - content: 'd', - filePath: 'd.java', - score: 3, - } - const chunkE: crossFile.CodeWhispererSupplementalContextItem = { - content: 'e', - filePath: 'e.java', - score: 4, - } - const chunkF: crossFile.CodeWhispererSupplementalContextItem = { - content: 'f', - filePath: 'f.java', - score: 5, - } - const chunkG: crossFile.CodeWhispererSupplementalContextItem = { - content: 'g', - filePath: 'g.java', - score: 6, - } - const chunks = [chunkA, chunkB, chunkC, chunkD, chunkE, chunkF, chunkG] - - assert.strictEqual(chunks.length, 7) - - const supplementalContext: CodeWhispererSupplementalContext = { - isUtg: false, - isProcessTimeout: false, - supplementalContextItems: chunks, - contentsLength: 25000, - latency: 0, - strategy: 'codemap', - } - - const actual = crossFile.truncateSuppelementalContext(supplementalContext) - assert.strictEqual(actual.supplementalContextItems.length, 5) - }) - - describe('truncate line by line', function () { - it('should return empty if empty string is provided', function () { - const input = '' - const actual = crossFile.truncateLineByLine(input, 50) - assert.strictEqual(actual, '') - }) - - it('should return empty if 0 max length is provided', function () { - const input = 'aaaaa' - const actual = crossFile.truncateLineByLine(input, 0) - assert.strictEqual(actual, '') - }) - - it('should flip the value if negative max length is provided', function () { - const input = `aaaaa${newLine}bbbbb` - const actual = crossFile.truncateLineByLine(input, -6) - const expected = crossFile.truncateLineByLine(input, 6) - assert.strictEqual(actual, expected) - assert.strictEqual(actual, 'aaaaa') - }) - }) - }) -}) - -function repeatString(s: string, n: number): string { - let output = '' - for (let i = 0; i < n; i++) { - output += s - } - - return output -} diff --git a/packages/amazonq/test/unit/codewhisperer/util/utgUtils.test.ts b/packages/amazonq/test/unit/codewhisperer/util/utgUtils.test.ts deleted file mode 100644 index 67359b8a6fc..00000000000 --- a/packages/amazonq/test/unit/codewhisperer/util/utgUtils.test.ts +++ /dev/null @@ -1,63 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as utgUtils from 'aws-core-vscode/codewhisperer' - -describe('shouldFetchUtgContext', () => { - it('fully supported language', function () { - assert.ok(utgUtils.shouldFetchUtgContext('java')) - }) - - it('partially supported language', () => { - assert.strictEqual(utgUtils.shouldFetchUtgContext('python'), false) - }) - - it('not supported language', () => { - assert.strictEqual(utgUtils.shouldFetchUtgContext('typescript'), undefined) - - assert.strictEqual(utgUtils.shouldFetchUtgContext('javascript'), undefined) - - assert.strictEqual(utgUtils.shouldFetchUtgContext('javascriptreact'), undefined) - - assert.strictEqual(utgUtils.shouldFetchUtgContext('typescriptreact'), undefined) - - assert.strictEqual(utgUtils.shouldFetchUtgContext('scala'), undefined) - - assert.strictEqual(utgUtils.shouldFetchUtgContext('shellscript'), undefined) - - assert.strictEqual(utgUtils.shouldFetchUtgContext('csharp'), undefined) - - assert.strictEqual(utgUtils.shouldFetchUtgContext('c'), undefined) - }) -}) - -describe('guessSrcFileName', function () { - it('should return undefined if no matching regex', function () { - assert.strictEqual(utgUtils.guessSrcFileName('Foo.java', 'java'), undefined) - assert.strictEqual(utgUtils.guessSrcFileName('folder1/foo.py', 'python'), undefined) - assert.strictEqual(utgUtils.guessSrcFileName('Bar.js', 'javascript'), undefined) - }) - - it('java', function () { - assert.strictEqual(utgUtils.guessSrcFileName('FooTest.java', 'java'), 'Foo.java') - assert.strictEqual(utgUtils.guessSrcFileName('FooTests.java', 'java'), 'Foo.java') - }) - - it('python', function () { - assert.strictEqual(utgUtils.guessSrcFileName('test_foo.py', 'python'), 'foo.py') - assert.strictEqual(utgUtils.guessSrcFileName('foo_test.py', 'python'), 'foo.py') - }) - - it('typescript', function () { - assert.strictEqual(utgUtils.guessSrcFileName('Foo.test.ts', 'typescript'), 'Foo.ts') - assert.strictEqual(utgUtils.guessSrcFileName('Foo.spec.ts', 'typescript'), 'Foo.ts') - }) - - it('javascript', function () { - assert.strictEqual(utgUtils.guessSrcFileName('Foo.test.js', 'javascript'), 'Foo.js') - assert.strictEqual(utgUtils.guessSrcFileName('Foo.spec.js', 'javascript'), 'Foo.js') - }) -}) diff --git a/packages/core/src/amazonq/commons/connector/baseMessenger.ts b/packages/core/src/amazonq/commons/connector/baseMessenger.ts index c26834c6fff..ea368965ea1 100644 --- a/packages/core/src/amazonq/commons/connector/baseMessenger.ts +++ b/packages/core/src/amazonq/commons/connector/baseMessenger.ts @@ -26,7 +26,7 @@ import { } from './connectorMessages' import { DeletedFileInfo, FollowUpTypes, NewFileInfo } from '../types' import { messageWithConversationId } from '../../../amazonqFeatureDev/userFacingText' -import { FeatureAuthState } from '../../../codewhisperer/util/authUtil' +import { auth2 } from 'aws-core-vscode/auth' export class Messenger { public constructor( @@ -191,19 +191,15 @@ export class Messenger { ) } - public async sendAuthNeededExceptionMessage(credentialState: FeatureAuthState, tabID: string) { + public async sendAuthNeededExceptionMessage(credentialState: auth2.AuthState, tabID: string) { let authType: AuthFollowUpType = 'full-auth' let message = AuthMessageDataMap[authType].message - switch (credentialState.amazonQ) { - case 'disconnected': + switch (credentialState) { + case 'notConnected': authType = 'full-auth' message = AuthMessageDataMap[authType].message break - case 'unsupported': - authType = 'use-supported-auth' - message = AuthMessageDataMap[authType].message - break case 'expired': authType = 're-auth' message = AuthMessageDataMap[authType].message diff --git a/packages/core/src/amazonq/explorer/amazonQTreeNode.ts b/packages/core/src/amazonq/explorer/amazonQTreeNode.ts index bbfd1bc1ff2..9b28fa80fe3 100644 --- a/packages/core/src/amazonq/explorer/amazonQTreeNode.ts +++ b/packages/core/src/amazonq/explorer/amazonQTreeNode.ts @@ -5,9 +5,9 @@ import * as vscode from 'vscode' import { ResourceTreeDataProvider, TreeNode } from '../../shared/treeview/resourceTreeDataProvider' -import { AuthState } from '../../codewhisperer/util/authUtil' import { createLearnMoreNode, createInstallQNode, createDismissNode } from './amazonQChildrenNodes' import { Commands } from '../../shared/vscode/commands2' +import { auth2 } from 'aws-core-vscode/auth' export class AmazonQNode implements TreeNode { public readonly id = 'amazonq' @@ -19,7 +19,7 @@ export class AmazonQNode implements TreeNode { private readonly onDidChangeVisibilityEmitter = new vscode.EventEmitter() public readonly onDidChangeVisibility = this.onDidChangeVisibilityEmitter.event - public static amazonQState: AuthState + public static amazonQState: auth2.AuthState private constructor() {} diff --git a/packages/core/src/amazonq/extApi.ts b/packages/core/src/amazonq/extApi.ts index 2eb16e4cde2..af98b73e59a 100644 --- a/packages/core/src/amazonq/extApi.ts +++ b/packages/core/src/amazonq/extApi.ts @@ -7,9 +7,14 @@ import vscode from 'vscode' import { VSCODE_EXTENSION_ID } from '../shared/extensions' import { SendMessageCommandOutput, SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client' import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseRequest } from '@amzn/codewhisperer-streaming' -import { FeatureAuthState } from '../codewhisperer/util/authUtil' +import { auth2 } from 'aws-core-vscode/auth' import { ToolkitError } from '../shared/errors' +/** + * @deprecated, for backwards comaptibility only. + */ +type OldAuthState = 'disconnected' | 'expired' | 'connected' + /** * This interface is used and exported by the amazon q extension. If you make a change here then * update the corresponding api implementation in packages/amazonq/src/api.ts @@ -21,7 +26,15 @@ export interface api { } authApi: { reauthIfNeeded(): Promise - getChatAuthState(): Promise + /** + * @deprecated, for backwards comaptibility only. + */ + getChatAuthState(): Promise<{ + codewhispererCore: OldAuthState + codewhispererChat: OldAuthState + amazonQ: OldAuthState + }> + getAuthState(): auth2.AuthState } } diff --git a/packages/core/src/amazonq/index.ts b/packages/core/src/amazonq/index.ts index 14c0e4a59a0..3b7737b3547 100644 --- a/packages/core/src/amazonq/index.ts +++ b/packages/core/src/amazonq/index.ts @@ -15,9 +15,6 @@ export { walkthroughInlineSuggestionsExample, walkthroughSecurityScanExample, } from './onboardingPage/walkthrough' -export { LspController } from './lsp/lspController' -export { LspClient } from './lsp/lspClient' -export * as lspClient from './lsp/lspClient' export { api } from './extApi' export { AmazonQChatViewProvider } from './webview/webView' export { amazonQHelpUrl } from '../shared/constants' @@ -40,8 +37,6 @@ export { ExtensionMessage } from '../amazonq/webview/ui/commands' export { CodeReference } from '../codewhispererChat/view/connector/connector' export { extractAuthFollowUp } from './util/authUtils' export { Messenger } from './commons/connector/baseMessenger' -export * from './lsp/config' -export * as WorkspaceLspInstaller from './lsp/workspaceInstaller' export * as secondaryAuth from '../auth/secondaryAuth' export * as authConnection from '../auth/connection' export * as featureConfig from './webview/generators/featureConfig' diff --git a/packages/core/src/amazonq/lsp/config.ts b/packages/core/src/amazonq/lsp/config.ts deleted file mode 100644 index 5670d0d0ce4..00000000000 --- a/packages/core/src/amazonq/lsp/config.ts +++ /dev/null @@ -1,31 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import { DevSettings } from '../../shared/settings' -import { getServiceEnvVarConfig } from '../../shared/vscode/env' - -export interface LspConfig { - manifestUrl: string - supportedVersions: string - id: string - suppressPromptPrefix: string - path?: string -} - -export const defaultAmazonQWorkspaceLspConfig: LspConfig = { - manifestUrl: 'https://aws-toolkit-language-servers.amazonaws.com/q-context/manifest.json', - supportedVersions: '0.1.47', - id: 'AmazonQ-Workspace', // used across IDEs for identifying global storage/local disk locations. Do not change. - suppressPromptPrefix: 'amazonQWorkspace', - path: undefined, -} - -export function getAmazonQWorkspaceLspConfig(): LspConfig { - return { - ...defaultAmazonQWorkspaceLspConfig, - ...(DevSettings.instance.getServiceConfig('amazonqWorkspaceLsp', {}) as LspConfig), - ...getServiceEnvVarConfig('amazonqWorkspaceLsp', Object.keys(defaultAmazonQWorkspaceLspConfig)), - } -} diff --git a/packages/core/src/amazonq/lsp/lspClient.ts b/packages/core/src/amazonq/lsp/lspClient.ts deleted file mode 100644 index eba89c961c4..00000000000 --- a/packages/core/src/amazonq/lsp/lspClient.ts +++ /dev/null @@ -1,378 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -/*! - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT License. See License.txt in the project root for license information. - */ -import * as vscode from 'vscode' -import { oneMB } from '../../shared/utilities/processUtils' -import * as path from 'path' -import * as nls from 'vscode-nls' -import * as crypto from 'crypto' -import * as jose from 'jose' - -import { Disposable, ExtensionContext } from 'vscode' - -import { LanguageClient, LanguageClientOptions } from 'vscode-languageclient' -import { - BuildIndexRequestPayload, - BuildIndexRequestType, - GetUsageRequestType, - IndexConfig, - QueryInlineProjectContextRequestType, - QueryVectorIndexRequestType, - UpdateIndexV2RequestPayload, - UpdateIndexV2RequestType, - QueryRepomapIndexRequestType, - GetRepomapIndexJSONRequestType, - Usage, - GetContextCommandItemsRequestType, - ContextCommandItem, - GetIndexSequenceNumberRequestType, - GetContextCommandPromptRequestType, - AdditionalContextPrompt, -} from './types' -import { CodeWhispererSettings } from '../../codewhisperer/util/codewhispererSettings' -import { fs } from '../../shared/fs/fs' -import { getLogger } from '../../shared/logger/logger' -import globals from '../../shared/extensionGlobals' -import { ResourcePaths } from '../../shared/lsp/types' -import { createServerOptions, validateNodeExe } from '../../shared/lsp/utils/platform' -import { waitUntil } from '../../shared/utilities/timeoutUtils' - -const localize = nls.loadMessageBundle() - -const key = crypto.randomBytes(32) -const logger = getLogger('amazonqWorkspaceLsp') - -/** - * LspClient manages the API call between VS Code extension and LSP server - * It encryptes the payload of API call. - */ -export class LspClient { - static #instance: LspClient - client: LanguageClient | undefined - - public static get instance() { - return (this.#instance ??= new this()) - } - - constructor() { - this.client = undefined - } - - async encrypt(payload: string) { - return await new jose.CompactEncrypt(new TextEncoder().encode(payload)) - .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' }) - .encrypt(key) - } - - async buildIndex(paths: string[], rootPath: string, config: IndexConfig) { - const payload: BuildIndexRequestPayload = { - filePaths: paths, - projectRoot: rootPath, - config: config, - language: '', - } - try { - const encryptedRequest = await this.encrypt(JSON.stringify(payload)) - const resp = await this.client?.sendRequest(BuildIndexRequestType, encryptedRequest) - return resp - } catch (e) { - logger.error(`buildIndex error: ${e}`) - return undefined - } - } - - async queryVectorIndex(request: string) { - try { - const encryptedRequest = await this.encrypt( - JSON.stringify({ - query: request, - }) - ) - const resp = await this.client?.sendRequest(QueryVectorIndexRequestType, encryptedRequest) - return resp - } catch (e) { - logger.error(`queryVectorIndex error: ${e}`) - return [] - } - } - - async queryInlineProjectContext(query: string, path: string, target: 'default' | 'codemap' | 'bm25') { - try { - const request = JSON.stringify({ - query: query, - filePath: path, - target, - }) - const encrypted = await this.encrypt(request) - const resp: any = await this.client?.sendRequest(QueryInlineProjectContextRequestType, encrypted) - return resp - } catch (e) { - logger.error(`queryInlineProjectContext error: ${e}`) - throw e - } - } - - async getLspServerUsage(): Promise { - if (this.client) { - return (await this.client.sendRequest(GetUsageRequestType, '')) as Usage - } - } - - async updateIndex(filePath: string[], mode: 'update' | 'remove' | 'add' | 'context_command_symbol_update') { - const payload: UpdateIndexV2RequestPayload = { - filePaths: filePath, - updateMode: mode, - } - try { - const encryptedRequest = await this.encrypt(JSON.stringify(payload)) - const resp = await this.client?.sendRequest(UpdateIndexV2RequestType, encryptedRequest) - return resp - } catch (e) { - logger.error(`updateIndex error: ${e}`) - return undefined - } - } - async queryRepomapIndex(filePaths: string[]) { - try { - const request = JSON.stringify({ - filePaths: filePaths, - }) - const resp: any = await this.client?.sendRequest(QueryRepomapIndexRequestType, await this.encrypt(request)) - return resp - } catch (e) { - logger.error(`QueryRepomapIndex error: ${e}`) - throw e - } - } - async getRepoMapJSON() { - try { - const request = JSON.stringify({}) - const resp: any = await this.client?.sendRequest( - GetRepomapIndexJSONRequestType, - await this.encrypt(request) - ) - return resp - } catch (e) { - logger.error(`queryInlineProjectContext error: ${e}`) - throw e - } - } - - async getContextCommandItems(): Promise { - try { - const workspaceFolders = vscode.workspace.workspaceFolders || [] - const request = JSON.stringify({ - workspaceFolders: workspaceFolders.map((it) => it.uri.fsPath), - }) - const resp: any = await this.client?.sendRequest( - GetContextCommandItemsRequestType, - await this.encrypt(request) - ) - return resp - } catch (e) { - logger.error(`getContextCommandItems error: ${e}`) - throw e - } - } - - async getContextCommandPrompt(contextCommandItems: ContextCommandItem[]): Promise { - try { - const request = JSON.stringify({ - contextCommands: contextCommandItems, - }) - const resp: any = await this.client?.sendRequest( - GetContextCommandPromptRequestType, - await this.encrypt(request) - ) - return resp || [] - } catch (e) { - logger.error(`getContextCommandPrompt error: ${e}`) - throw e - } - } - - async getIndexSequenceNumber(): Promise { - try { - const request = JSON.stringify({}) - const resp: any = await this.client?.sendRequest( - GetIndexSequenceNumberRequestType, - await this.encrypt(request) - ) - return resp - } catch (e) { - logger.error(`getIndexSequenceNumber error: ${e}`) - throw e - } - } - - async waitUntilReady() { - return waitUntil( - async () => { - if (this.client === undefined) { - return false - } - await this.client.onReady() - return true - }, - { interval: 500, timeout: 60_000 * 3, truthy: true } - ) - } -} - -/** - * Activates the language server (assumes the LSP server has already been downloaded): - * 1. start LSP server running over IPC protocol. - * 2. create a output channel named Amazon Q Language Server. - */ -export async function activate(extensionContext: ExtensionContext, resourcePaths: ResourcePaths) { - LspClient.instance // Tickle the singleton... :/ - const toDispose = extensionContext.subscriptions - - let rangeFormatting: Disposable | undefined - // The debug options for the server - // --inspect=6009: runs the server in Node's Inspector mode so VS Code can attach to the server for debugging - const debugOptions = { execArgv: ['--nolazy', '--preserve-symlinks', '--stdio'] } - const workerThreads = CodeWhispererSettings.instance.getIndexWorkerThreads() - const gpu = CodeWhispererSettings.instance.isLocalIndexGPUEnabled() - - if (gpu) { - process.env.Q_ENABLE_GPU = 'true' - } else { - delete process.env.Q_ENABLE_GPU - } - if (workerThreads > 0 && workerThreads < 100) { - process.env.Q_WORKER_THREADS = workerThreads.toString() - } else { - delete process.env.Q_WORKER_THREADS - } - - const serverModule = resourcePaths.lsp - const memoryWarnThreshold = 800 * oneMB - - const serverOptions = createServerOptions({ - encryptionKey: key, - executable: [resourcePaths.node], - serverModule, - // TODO(jmkeyes): we always use the debug options...? - execArgv: debugOptions.execArgv, - warnThresholds: { memory: memoryWarnThreshold }, - }) - - const documentSelector = [{ scheme: 'file', language: '*' }] - - await validateNodeExe([resourcePaths.node], resourcePaths.lsp, debugOptions.execArgv, logger) - - // Options to control the language client - const clientOptions: LanguageClientOptions = { - // Register the server for json documents - documentSelector, - initializationOptions: { - handledSchemaProtocols: ['file', 'untitled'], // language server only loads file-URI. Fetching schemas with other protocols ('http'...) are made on the client. - provideFormatter: false, // tell the server to not provide formatting capability and ignore the `aws.stepfunctions.asl.format.enable` setting. - // this is used by LSP to determine index cache path, move to this folder so that when extension updates index is not deleted. - extensionPath: path.join(fs.getUserHomeDir(), '.aws', 'amazonq', 'cache'), - }, - // Log to the Amazon Q Logs so everything is in a single channel - // TODO: Add prefix to the language server logs so it is easier to search - outputChannel: globals.logOutputChannel, - } - - // Create the language client and start the client. - LspClient.instance.client = new LanguageClient( - 'amazonq', - localize('amazonq.server.name', 'Amazon Q Language Server'), - serverOptions, - clientOptions - ) - LspClient.instance.client.registerProposedFeatures() - - const disposable = LspClient.instance.client.start() - toDispose.push(disposable) - - let savedDocument: vscode.Uri | undefined = undefined - - const onAdd = async (filePaths: string[]) => { - const indexSeqNum = await LspClient.instance.getIndexSequenceNumber() - await LspClient.instance.updateIndex(filePaths, 'add') - await waitUntil( - async () => { - const newIndexSeqNum = await LspClient.instance.getIndexSequenceNumber() - if (newIndexSeqNum > indexSeqNum) { - await vscode.commands.executeCommand(`aws.amazonq.updateContextCommandItems`) - return true - } - return false - }, - { interval: 500, timeout: 5_000, truthy: true } - ) - } - const onRemove = async (filePaths: string[]) => { - const indexSeqNum = await LspClient.instance.getIndexSequenceNumber() - await LspClient.instance.updateIndex(filePaths, 'remove') - await waitUntil( - async () => { - const newIndexSeqNum = await LspClient.instance.getIndexSequenceNumber() - if (newIndexSeqNum > indexSeqNum) { - await vscode.commands.executeCommand(`aws.amazonq.updateContextCommandItems`) - return true - } - return false - }, - { interval: 500, timeout: 5_000, truthy: true } - ) - } - - toDispose.push( - vscode.workspace.onDidSaveTextDocument((document) => { - if (document.uri.scheme !== 'file') { - return - } - savedDocument = document.uri - }), - vscode.window.onDidChangeActiveTextEditor((editor) => { - if (savedDocument && editor && editor.document.uri.fsPath !== savedDocument.fsPath) { - void LspClient.instance.updateIndex([savedDocument.fsPath], 'update') - } - // user created a new empty file using File -> New File - // these events will not be captured by vscode.workspace.onDidCreateFiles - // because it was created by File Explorer(Win) or Finder(MacOS) - // TODO: consider using a high performance fs watcher - if (editor?.document.getText().length === 0) { - void onAdd([editor.document.uri.fsPath]) - } - }), - vscode.workspace.onDidCreateFiles(async (e) => { - await onAdd(e.files.map((f) => f.fsPath)) - }), - vscode.workspace.onDidDeleteFiles(async (e) => { - await onRemove(e.files.map((f) => f.fsPath)) - }), - vscode.workspace.onDidRenameFiles(async (e) => { - await onRemove(e.files.map((f) => f.oldUri.fsPath)) - await onAdd(e.files.map((f) => f.newUri.fsPath)) - }) - ) - - return LspClient.instance.client.onReady().then( - () => { - const disposableFunc = { dispose: () => rangeFormatting?.dispose() as void } - toDispose.push(disposableFunc) - }, - (reason) => { - logger.error('client.onReady() failed: %O', reason) - } - ) -} - -export async function deactivate(): Promise { - if (!LspClient.instance.client) { - return undefined - } - return LspClient.instance.client.stop() -} diff --git a/packages/core/src/amazonq/lsp/lspController.ts b/packages/core/src/amazonq/lsp/lspController.ts deleted file mode 100644 index 5a1b84b7c49..00000000000 --- a/packages/core/src/amazonq/lsp/lspController.ts +++ /dev/null @@ -1,237 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import * as path from 'path' -import { getLogger } from '../../shared/logger/logger' -import { CurrentWsFolders, collectFilesForIndex } from '../../shared/utilities/workspaceUtils' -import { activate as activateLsp, LspClient } from './lspClient' -import { telemetry } from '../../shared/telemetry/telemetry' -import { isCloud9 } from '../../shared/extensionUtilities' -import globals, { isWeb } from '../../shared/extensionGlobals' -import { isAmazonLinux2 } from '../../shared/vscode/env' -import { WorkspaceLspInstaller } from './workspaceInstaller' -import { lspSetupStage } from '../../shared/lsp/utils/setupStage' -import { RelevantTextDocumentAddition } from '../../codewhispererChat/controllers/chat/model' -import { waitUntil } from '../../shared/utilities/timeoutUtils' - -export interface Chunk { - readonly filePath: string - readonly content: string - readonly context?: string - readonly relativePath?: string - readonly programmingLanguage?: string - readonly startLine?: number - readonly endLine?: number -} -export interface BuildIndexConfig { - startUrl?: string - maxIndexSize: number - isVectorIndexEnabled: boolean -} - -/* - * LSP Controller manages the status of Amazon Q Workspace Indexing LSP: - * 1. Downloading, verifying and installing LSP using DEXP LSP manifest and CDN. - * 2. Managing the LSP states. There are a couple of possible LSP states: - * Not installed. Installed. Running. Indexing. Indexing Done. - * LSP Controller converts the input and output of LSP APIs. - * The IDE extension code should invoke LSP API via this controller. - * 3. It perform pre-process and post process of LSP APIs - * Pre-process the input to Index Files API - * Post-process the output from Query API - */ -export class LspController { - static #instance: LspController - private _isIndexingInProgress = false - private _contextCommandSymbolsUpdated = false - private logger = getLogger('amazonqWorkspaceLsp') - - public static get instance() { - return (this.#instance ??= new this()) - } - - isIndexingInProgress() { - return this._isIndexingInProgress - } - - async query(s: string): Promise { - const chunks: Chunk[] | undefined = await LspClient.instance.queryVectorIndex(s) - const resp: RelevantTextDocumentAddition[] = [] - if (chunks) { - for (const chunk of chunks) { - const text = chunk.context ? chunk.context : chunk.content - if (chunk.programmingLanguage && chunk.programmingLanguage !== 'unknown') { - resp.push({ - text: text, - relativeFilePath: chunk.relativePath ? chunk.relativePath : path.basename(chunk.filePath), - programmingLanguage: { - languageName: chunk.programmingLanguage, - }, - startLine: chunk.startLine ?? -1, - endLine: chunk.endLine ?? -1, - }) - } else { - resp.push({ - text: text, - relativeFilePath: chunk.relativePath ? chunk.relativePath : path.basename(chunk.filePath), - startLine: chunk.startLine ?? -1, - endLine: chunk.endLine ?? -1, - }) - } - } - } - return resp - } - - async queryInlineProjectContext(query: string, path: string, target: 'bm25' | 'codemap' | 'default') { - try { - return await LspClient.instance.queryInlineProjectContext(query, path, target) - } catch (e) { - if (e instanceof Error) { - this.logger.error(`unexpected error while querying inline project context, e=${e.message}`) - } - return [] - } - } - - async buildIndex(buildIndexConfig: BuildIndexConfig) { - this.logger.info(`Starting to build index of project`) - const start = performance.now() - const projPaths = (vscode.workspace.workspaceFolders ?? []).map((folder) => folder.uri.fsPath) - if (projPaths.length === 0) { - this.logger.info(`Skipping building index. No projects found in workspace`) - return - } - projPaths.sort() - try { - this._isIndexingInProgress = true - const projRoot = projPaths[0] - const files = await collectFilesForIndex( - projPaths, - vscode.workspace.workspaceFolders as CurrentWsFolders, - true, - buildIndexConfig.maxIndexSize * 1024 * 1024 - ) - const totalSizeBytes = files.reduce( - (accumulator, currentFile) => accumulator + currentFile.fileSizeBytes, - 0 - ) - this.logger.info(`Found ${files.length} files in current project ${projPaths}`) - const config = buildIndexConfig.isVectorIndexEnabled ? 'all' : 'default' - const r = files.map((f) => f.fileUri.fsPath) - const resp = await LspClient.instance.buildIndex(r, projRoot, config) - if (resp) { - this.logger.debug(`Finish building index of project`) - const usage = await LspClient.instance.getLspServerUsage() - telemetry.amazonq_indexWorkspace.emit({ - duration: performance.now() - start, - result: 'Succeeded', - amazonqIndexFileCount: files.length, - amazonqIndexMemoryUsageInMB: usage ? usage.memoryUsage / (1024 * 1024) : undefined, - amazonqIndexCpuUsagePercentage: usage ? usage.cpuUsage : undefined, - amazonqIndexFileSizeInMB: totalSizeBytes / (1024 * 1024), - amazonqVectorIndexEnabled: buildIndexConfig.isVectorIndexEnabled, - credentialStartUrl: buildIndexConfig.startUrl, - }) - } else { - this.logger.error(`Failed to build index of project`) - telemetry.amazonq_indexWorkspace.emit({ - duration: performance.now() - start, - result: 'Failed', - amazonqIndexFileCount: 0, - amazonqIndexFileSizeInMB: 0, - amazonqVectorIndexEnabled: buildIndexConfig.isVectorIndexEnabled, - reason: `Unknown`, - }) - } - } catch (error) { - // TODO: use telemetry.run() - this.logger.error(`Failed to build index of project`) - telemetry.amazonq_indexWorkspace.emit({ - duration: performance.now() - start, - result: 'Failed', - amazonqIndexFileCount: 0, - amazonqIndexFileSizeInMB: 0, - amazonqVectorIndexEnabled: buildIndexConfig.isVectorIndexEnabled, - reason: `${error instanceof Error ? error.name : 'Unknown'}`, - reasonDesc: `Error when building index. ${error instanceof Error ? error.message : error}`, - }) - } finally { - this._isIndexingInProgress = false - } - } - - async trySetupLsp(context: vscode.ExtensionContext, buildIndexConfig: BuildIndexConfig) { - if (isCloud9() || isWeb() || isAmazonLinux2()) { - this.logger.warn('Skipping LSP setup. LSP is not compatible with the current environment. ') - // do not do anything if in Cloud9 or Web mode or in AL2 (AL2 does not support node v18+) - return - } - setImmediate(async () => { - try { - await this.setupLsp(context) - await vscode.commands.executeCommand(`aws.amazonq.updateContextCommandItems`) - void LspController.instance.buildIndex(buildIndexConfig) - // log the LSP server CPU and Memory usage per 30 minutes. - globals.clock.setInterval( - async () => { - const usage = await LspClient.instance.getLspServerUsage() - if (usage) { - this.logger.info( - `LSP server CPU ${usage.cpuUsage}%, LSP server Memory ${ - usage.memoryUsage / (1024 * 1024) - }MB ` - ) - } - }, - 30 * 60 * 1000 - ) - } catch (e) { - this.logger.error(`LSP failed to activate ${e}`) - } - }) - } - /** - * Updates context command symbols once per session by synchronizing with the LSP client index. - * Context menu will contain file and folders to begin with, - * then this asynchronous function should be invoked after the files and folders are found - * the LSP then further starts to parse workspace and find symbols, which takes - * anywhere from 5 seconds to about 40 seconds, depending on project size. - * @returns {Promise} - */ - async updateContextCommandSymbolsOnce() { - if (this._contextCommandSymbolsUpdated) { - return - } - this._contextCommandSymbolsUpdated = true - getLogger().debug(`Start adding symbols to context picker menu`) - try { - const indexSeqNum = await LspClient.instance.getIndexSequenceNumber() - await LspClient.instance.updateIndex([], 'context_command_symbol_update') - await waitUntil( - async () => { - const newIndexSeqNum = await LspClient.instance.getIndexSequenceNumber() - if (newIndexSeqNum > indexSeqNum) { - await vscode.commands.executeCommand(`aws.amazonq.updateContextCommandItems`) - return true - } - return false - }, - { interval: 1000, timeout: 60_000, truthy: true } - ) - } catch (err) { - this.logger.error(`Failed to find symbols`) - } - } - - private async setupLsp(context: vscode.ExtensionContext) { - await lspSetupStage('all', async () => { - const installResult = await new WorkspaceLspInstaller().resolve() - await lspSetupStage('launch', async () => activateLsp(context, installResult.resourcePaths)) - this.logger.info('LSP activated') - }) - } -} diff --git a/packages/core/src/amazonq/lsp/types.ts b/packages/core/src/amazonq/lsp/types.ts deleted file mode 100644 index 2940ce240c8..00000000000 --- a/packages/core/src/amazonq/lsp/types.ts +++ /dev/null @@ -1,150 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import { RequestType } from 'vscode-languageserver' - -export type IndexRequestPayload = { - filePaths: string[] - rootPath: string - refresh: boolean -} - -export type ClearRequest = string - -export const ClearRequestType: RequestType = new RequestType('lsp/clear') - -export type QueryRequest = string - -export const QueryRequestType: RequestType = new RequestType('lsp/query') - -export type GetUsageRequest = string - -export const GetUsageRequestType: RequestType = new RequestType('lsp/getUsage') - -export interface Usage { - memoryUsage: number - cpuUsage: number -} - -export type BuildIndexRequestPayload = { - filePaths: string[] - projectRoot: string - config: string - language: string -} - -export type BuildIndexRequest = string - -export const BuildIndexRequestType: RequestType = new RequestType('lsp/buildIndex') - -export type UpdateIndexV2Request = string - -export type UpdateIndexV2RequestPayload = { filePaths: string[]; updateMode: string } - -export const UpdateIndexV2RequestType: RequestType = new RequestType( - 'lsp/updateIndexV2' -) - -export type QueryInlineProjectContextRequest = string -export type QueryInlineProjectContextRequestPayload = { - query: string - filePath: string - target: string -} -export const QueryInlineProjectContextRequestType: RequestType = - new RequestType('lsp/queryInlineProjectContext') - -export type QueryVectorIndexRequestPayload = { query: string } - -export type QueryVectorIndexRequest = string - -export const QueryVectorIndexRequestType: RequestType = new RequestType( - 'lsp/queryVectorIndex' -) - -export type IndexConfig = 'all' | 'default' - -// RepoMapData -export type QueryRepomapIndexRequestPayload = { filePaths: string[] } -export type QueryRepomapIndexRequest = string -export const QueryRepomapIndexRequestType: RequestType = new RequestType( - 'lsp/queryRepomapIndex' -) -export type GetRepomapIndexJSONRequest = string -export const GetRepomapIndexJSONRequestType: RequestType = new RequestType( - 'lsp/getRepomapIndexJSON' -) - -export type GetContextCommandItemsRequestPayload = { workspaceFolders: string[] } -export type GetContextCommandItemsRequest = string -export const GetContextCommandItemsRequestType: RequestType = new RequestType( - 'lsp/getContextCommandItems' -) - -export type GetIndexSequenceNumberRequest = string -export const GetIndexSequenceNumberRequestType: RequestType = new RequestType( - 'lsp/getIndexSequenceNumber' -) - -export type ContextCommandItemType = 'file' | 'folder' | 'code' - -export type SymbolType = - | 'Class' - | 'Function' - | 'Interface' - | 'Type' - | 'Enum' - | 'Struct' - | 'Delegate' - | 'Namespace' - | 'Object' - | 'Module' - | 'Method' - -export interface Position { - line: number - column: number -} -export interface Span { - start: Position - end: Position -} - -// LSP definition of DocumentSymbol - -export interface DocumentSymbol { - name: string - kind: SymbolType - range: Span -} - -export interface ContextCommandItem { - workspaceFolder: string - type: ContextCommandItemType - relativePath: string - symbol?: DocumentSymbol - id?: string -} - -export type GetContextCommandPromptRequestPayload = { - contextCommands: { - workspaceFolder: string - type: 'file' | 'folder' - relativePath: string - }[] -} -export type GetContextCommandPromptRequest = string -export const GetContextCommandPromptRequestType: RequestType = - new RequestType('lsp/getContextCommandPrompt') - -export interface AdditionalContextPrompt { - content: string - name: string - description: string - startLine: number - endLine: number - filePath: string - relativePath: string -} diff --git a/packages/core/src/amazonq/lsp/workspaceInstaller.ts b/packages/core/src/amazonq/lsp/workspaceInstaller.ts deleted file mode 100644 index 99e70f20cbf..00000000000 --- a/packages/core/src/amazonq/lsp/workspaceInstaller.ts +++ /dev/null @@ -1,39 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import path from 'path' -import { ResourcePaths } from '../../shared/lsp/types' -import { getNodeExecutableName } from '../../shared/lsp/utils/platform' -import { fs } from '../../shared/fs/fs' -import { BaseLspInstaller } from '../../shared/lsp/baseLspInstaller' -import { getAmazonQWorkspaceLspConfig, LspConfig } from './config' - -export class WorkspaceLspInstaller extends BaseLspInstaller { - constructor(lspConfig: LspConfig = getAmazonQWorkspaceLspConfig()) { - super(lspConfig, 'amazonqWorkspaceLsp') - } - - protected override async postInstall(assetDirectory: string): Promise { - const resourcePaths = this.resourcePaths(assetDirectory) - await fs.chmod(resourcePaths.node, 0o755) - } - - protected override resourcePaths(assetDirectory?: string): ResourcePaths { - // local version - if (!assetDirectory) { - return { - lsp: this.config.path ?? '', - node: getNodeExecutableName(), - } - } - - const lspNodeName = - process.platform === 'win32' ? getNodeExecutableName() : `node-${process.platform}-${process.arch}` - return { - lsp: path.join(assetDirectory, `qserver-${process.platform}-${process.arch}/qserver/lspServer.js`), - node: path.join(assetDirectory, lspNodeName), - } - } -} diff --git a/packages/core/src/amazonq/session/sessionState.ts b/packages/core/src/amazonq/session/sessionState.ts index 1f206c23159..10e7eb67fdd 100644 --- a/packages/core/src/amazonq/session/sessionState.ts +++ b/packages/core/src/amazonq/session/sessionState.ts @@ -238,7 +238,7 @@ export abstract class BasePrepareCodeGenState implements SessionState { const uploadId = await telemetry.amazonq_createUpload.run(async (span) => { span.record({ amazonqConversationId: this.config.conversationId, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, }) const { zipFileBuffer, zipFileChecksum } = await this.prepareProjectZip( this.config.workspaceRoots, @@ -357,7 +357,7 @@ export abstract class BaseCodeGenState extends CodeGenBase implements SessionSta span.record({ amazonqConversationId: this.config.conversationId, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, }) action.telemetry.setGenerateCodeIteration(this.currentIteration) diff --git a/packages/core/src/amazonq/util/authUtils.ts b/packages/core/src/amazonq/util/authUtils.ts index e310dbe9823..f3e48a22cc0 100644 --- a/packages/core/src/amazonq/util/authUtils.ts +++ b/packages/core/src/amazonq/util/authUtils.ts @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { FeatureAuthState } from '../../codewhisperer/util/authUtil' +import { AuthState } from '../../auth/auth2' import { AuthFollowUpType, AuthMessageDataMap } from '../auth/model' /** @@ -15,20 +15,13 @@ import { AuthFollowUpType, AuthMessageDataMap } from '../auth/model' * - authType: The type of authentication follow-up required (AuthFollowUpType) * - message: The corresponding message for the determined auth type */ -export function extractAuthFollowUp(credentialState: FeatureAuthState) { +export function extractAuthFollowUp(credentialState: AuthState) { let authType: AuthFollowUpType = 'full-auth' let message = AuthMessageDataMap[authType].message - if (credentialState.codewhispererChat === 'disconnected' && credentialState.codewhispererCore === 'disconnected') { + if (credentialState === 'notConnected') { authType = 'full-auth' message = AuthMessageDataMap[authType].message - } - - if (credentialState.codewhispererCore === 'connected' && credentialState.codewhispererChat === 'expired') { - authType = 'missing_scopes' - message = AuthMessageDataMap[authType].message - } - - if (credentialState.codewhispererChat === 'expired' && credentialState.codewhispererCore === 'expired') { + } else if (credentialState === 'expired') { authType = 're-auth' message = AuthMessageDataMap[authType].message } diff --git a/packages/core/src/amazonq/util/files.ts b/packages/core/src/amazonq/util/files.ts index afa0b674928..fce56e37a08 100644 --- a/packages/core/src/amazonq/util/files.ts +++ b/packages/core/src/amazonq/util/files.ts @@ -278,7 +278,7 @@ export function registerNewFiles( telemetry.toolkit_trackScenario.emit({ count: 1, amazonqConversationId: conversationId, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, scenario: 'wsOrphanedDocuments', }) getLogger().error(`No workspace folder found for file: ${zipFilePath} and prefix: ${prefix}`) diff --git a/packages/core/src/amazonq/webview/generators/webViewContent.ts b/packages/core/src/amazonq/webview/generators/webViewContent.ts index bc637a8f6c6..ddbee4f0e9b 100644 --- a/packages/core/src/amazonq/webview/generators/webViewContent.ts +++ b/packages/core/src/amazonq/webview/generators/webViewContent.ts @@ -70,7 +70,7 @@ export class WebViewContentGenerator { : AuthUtil.instance.regionProfileManager.activeRegionProfile const regionProfileString: string = JSON.stringify(regionProfile) - const authState = (await AuthUtil.instance.getChatAuthState()).amazonQ + const authState = AuthUtil.instance.getAuthState() return ` diff --git a/packages/core/src/amazonqDoc/app.ts b/packages/core/src/amazonqDoc/app.ts index 929cf1d45de..d348583a798 100644 --- a/packages/core/src/amazonqDoc/app.ts +++ b/packages/core/src/amazonqDoc/app.ts @@ -81,7 +81,7 @@ export function init(appContext: AmazonQAppInitContext) { appContext.registerWebViewToAppMessagePublisher(new MessagePublisher(docChatUIInputEventEmitter), 'doc') const debouncedEvent = debounce(async () => { - const authenticated = (await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected' + const authenticated = AuthUtil.instance.getAuthState() === 'connected' let authenticatingSessionIDs: string[] = [] if (authenticated) { const authenticatingSessions = sessionStorage.getAuthenticatingSessions() @@ -97,9 +97,10 @@ export function init(appContext: AmazonQAppInitContext) { messenger.sendAuthenticationUpdate(authenticated, authenticatingSessionIDs) }, 500) - AuthUtil.instance.secondaryAuth.onDidChangeActiveConnection(() => { + AuthUtil.instance.onDidChangeConnectionState(() => { return debouncedEvent() }) + AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(() => { return debouncedEvent() }) diff --git a/packages/core/src/amazonqDoc/controllers/chat/controller.ts b/packages/core/src/amazonqDoc/controllers/chat/controller.ts index ab6045e75ce..390a959b937 100644 --- a/packages/core/src/amazonqDoc/controllers/chat/controller.ts +++ b/packages/core/src/amazonqDoc/controllers/chat/controller.ts @@ -232,9 +232,8 @@ export class DocController { const workspaceFolderName = vscode.workspace.workspaceFolders?.[0].name || '' - const authState = await AuthUtil.instance.getChatAuthState() - - if (authState.amazonQ !== 'connected') { + const authState = AuthUtil.instance.getAuthState() + if (authState !== 'connected') { await this.messenger.sendAuthNeededExceptionMessage(authState, data.tabID) session.isAuthenticating = true return @@ -465,8 +464,8 @@ export class DocController { try { getLogger().debug(`${featureName}: Processing message: ${message.message}`) - const authState = await AuthUtil.instance.getChatAuthState() - if (authState.amazonQ !== 'connected') { + const authState = AuthUtil.instance.getAuthState() + if (authState !== 'connected') { await this.messenger.sendAuthNeededExceptionMessage(authState, message.tabID) session.isAuthenticating = true return @@ -501,8 +500,8 @@ export class DocController { docGenerationTask.folderPath = '' docGenerationTask.mode = Mode.NONE - const authState = await AuthUtil.instance.getChatAuthState() - if (authState.amazonQ !== 'connected') { + const authState = AuthUtil.instance.getAuthState() + if (authState !== 'connected') { void this.messenger.sendAuthNeededExceptionMessage(authState, message.tabID) session.isAuthenticating = true return diff --git a/packages/core/src/amazonqDoc/session/session.ts b/packages/core/src/amazonqDoc/session/session.ts index e3eb29d6d32..8cfb5283626 100644 --- a/packages/core/src/amazonqDoc/session/session.ts +++ b/packages/core/src/amazonqDoc/session/session.ts @@ -91,7 +91,10 @@ export class Session { this._conversationId = await this.proxyClient.createConversation() getLogger().info(logWithConversationId(this.conversationId)) - span.record({ amazonqConversationId: this._conversationId, credentialStartUrl: AuthUtil.instance.startUrl }) + span.record({ + amazonqConversationId: this._conversationId, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, + }) }) this._state = new DocPrepareCodeGenState( diff --git a/packages/core/src/amazonqFeatureDev/app.ts b/packages/core/src/amazonqFeatureDev/app.ts index a016d2ba481..a3130ca4f74 100644 --- a/packages/core/src/amazonqFeatureDev/app.ts +++ b/packages/core/src/amazonqFeatureDev/app.ts @@ -87,7 +87,7 @@ export function init(appContext: AmazonQAppInitContext) { ) const debouncedEvent = debounce(async () => { - const authenticated = (await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected' + const authenticated = AuthUtil.instance.getAuthState() === 'connected' let authenticatingSessionIDs: string[] = [] if (authenticated) { const authenticatingSessions = sessionStorage.getAuthenticatingSessions() @@ -103,9 +103,10 @@ export function init(appContext: AmazonQAppInitContext) { messenger.sendAuthenticationUpdate(authenticated, authenticatingSessionIDs) }, 500) - AuthUtil.instance.secondaryAuth.onDidChangeActiveConnection(() => { + AuthUtil.instance.onDidChangeConnectionState(() => { return debouncedEvent() }) + AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(() => { return debouncedEvent() }) diff --git a/packages/core/src/amazonqFeatureDev/client/featureDev.ts b/packages/core/src/amazonqFeatureDev/client/featureDev.ts index 62e870b51fe..419f6969cc6 100644 --- a/packages/core/src/amazonqFeatureDev/client/featureDev.ts +++ b/packages/core/src/amazonqFeatureDev/client/featureDev.ts @@ -51,7 +51,7 @@ const writeAPIRetryOptions = { // Create a client for featureDev proxy client based off of aws sdk v2 export async function createFeatureDevProxyClient(options?: Partial): Promise { - const bearerToken = await AuthUtil.instance.getBearerToken() + const bearerToken = await AuthUtil.instance.getToken() const cwsprConfig = getCodewhispererConfig() return (await globals.sdkClientBuilder.createAwsService( Service, diff --git a/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts b/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts index bdf73eada07..4843b3d6793 100644 --- a/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts +++ b/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts @@ -216,14 +216,14 @@ export class FeatureDevController { amazonqConversationId: session?.conversationId, value: 1, result: 'Succeeded', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, }) } else if (vote === 'downvote') { telemetry.amazonq_codeGenerationThumbsDown.emit({ amazonqConversationId: session?.conversationId, value: 1, result: 'Succeeded', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, }) } } @@ -395,8 +395,8 @@ export class FeatureDevController { session.latestMessage = message.message await session.disableFileList() - const authState = await AuthUtil.instance.getChatAuthState() - if (authState.amazonQ !== 'connected') { + const authState = AuthUtil.instance.getAuthState() + if (authState !== 'connected') { await this.messenger.sendAuthNeededExceptionMessage(authState, message.tabID) session.isAuthenticating = true return @@ -717,7 +717,7 @@ export class FeatureDevController { amazonqConversationId: session.conversationId, enabled: true, result: 'Succeeded', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, }) // Unblock the message button this.messenger.sendAsyncEventProgress(message.tabID, false, undefined) @@ -828,7 +828,7 @@ export class FeatureDevController { } telemetry.amazonq_modifySourceFolder.emit({ - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, amazonqConversationId: session.conversationId, ...metricData, }) @@ -920,7 +920,7 @@ export class FeatureDevController { amazonqConversationId: session.conversationId, enabled: true, result: 'Succeeded', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, }) const workspacePrefixMapping = getWorkspaceFoldersByPrefixes(session.config.workspaceFolders) @@ -970,8 +970,8 @@ export class FeatureDevController { session = await this.sessionStorage.getSession(message.tabID) getLogger().debug(`${featureName}: Session created with id: ${session.tabID}`) - const authState = await AuthUtil.instance.getChatAuthState() - if (authState.amazonQ !== 'connected') { + const authState = AuthUtil.instance.getAuthState() + if (authState !== 'connected') { void this.messenger.sendAuthNeededExceptionMessage(authState, message.tabID) session.isAuthenticating = true return @@ -1078,7 +1078,7 @@ export class FeatureDevController { if (amazonqNumberOfFilesAccepted > 0 && !session.acceptCodeTelemetrySent) { session.updateAcceptCodeTelemetrySent(true) telemetry.amazonq_isAcceptedCodeChanges.emit({ - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, amazonqConversationId: session.conversationId, amazonqNumberOfFilesAccepted, enabled: true, diff --git a/packages/core/src/amazonqFeatureDev/session/session.ts b/packages/core/src/amazonqFeatureDev/session/session.ts index c1fc81a4701..f89db0aa9e7 100644 --- a/packages/core/src/amazonqFeatureDev/session/session.ts +++ b/packages/core/src/amazonqFeatureDev/session/session.ts @@ -92,7 +92,10 @@ export class Session { this._conversationId = await this.proxyClient.createConversation() getLogger().info(logWithConversationId(this.conversationId)) - span.record({ amazonqConversationId: this._conversationId, credentialStartUrl: AuthUtil.instance.startUrl }) + span.record({ + amazonqConversationId: this._conversationId, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, + }) }) this._state = new FeatureDevPrepareCodeGenState( diff --git a/packages/core/src/amazonqGumby/app.ts b/packages/core/src/amazonqGumby/app.ts index 21182b38155..022d2f58f30 100644 --- a/packages/core/src/amazonqGumby/app.ts +++ b/packages/core/src/amazonqGumby/app.ts @@ -49,7 +49,7 @@ export function init(appContext: AmazonQAppInitContext) { appContext.registerWebViewToAppMessagePublisher(new MessagePublisher(gumbyChatUIInputEventEmitter), 'gumby') const debouncedEvent = debounce(async () => { - const authenticated = (await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected' + const authenticated = AuthUtil.instance.getAuthState() === 'connected' let authenticatingSessionID = '' if (authenticated) { @@ -64,7 +64,7 @@ export function init(appContext: AmazonQAppInitContext) { messenger.sendAuthenticationUpdate(authenticated, [authenticatingSessionID]) }, 500) - AuthUtil.instance.secondaryAuth.onDidChangeActiveConnection(() => { + AuthUtil.instance.onDidChangeConnectionState(() => { return debouncedEvent() }) AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(() => { diff --git a/packages/core/src/amazonqGumby/chat/controller/controller.ts b/packages/core/src/amazonqGumby/chat/controller/controller.ts index 57367143cd4..6d1e94777fe 100644 --- a/packages/core/src/amazonqGumby/chat/controller/controller.ts +++ b/packages/core/src/amazonqGumby/chat/controller/controller.ts @@ -159,8 +159,8 @@ export class GumbyController { try { getLogger().debug(`${featureName}: Session created with id: ${session.tabID}`) - const authState = await AuthUtil.instance.getChatAuthState() - if (authState.amazonQ !== 'connected') { + const authState = AuthUtil.instance.getAuthState() + if (authState !== 'connected') { void this.messenger.sendAuthNeededExceptionMessage(authState, tabID) session.isAuthenticating = true return @@ -259,11 +259,11 @@ export class GumbyController { credentialSourceId: authType, }) - const authState = await AuthUtil.instance.getChatAuthState() - if (authState.amazonQ !== 'connected') { + const authState = AuthUtil.instance.getAuthState() + if (authState !== 'connected') { this.sessionStorage.getSession().isAuthenticating = true await this.messenger.sendAuthNeededExceptionMessage(authState, message.tabID) - throw new AuthError('Not connected to Amazon Q', `AuthState=${authState.amazonQ}`) + throw new AuthError('Not connected to Amazon Q', `AuthState=${authState}`) } this.messenger.sendTransformationIntroduction(message.tabID) }) @@ -542,8 +542,8 @@ export class GumbyController { this.messenger.sendCompilationFinished(tabID) // since compilation can potentially take a long time, double check auth - const authState = await AuthUtil.instance.getChatAuthState() - if (authState.amazonQ !== 'connected') { + const authState = AuthUtil.instance.getAuthState() + if (authState !== 'connected') { void this.messenger.sendAuthNeededExceptionMessage(authState, tabID) this.sessionStorage.getSession().isAuthenticating = true return diff --git a/packages/core/src/amazonqGumby/chat/controller/messenger/messenger.ts b/packages/core/src/amazonqGumby/chat/controller/messenger/messenger.ts index 9d15271aa1e..93ebb6120e7 100644 --- a/packages/core/src/amazonqGumby/chat/controller/messenger/messenger.ts +++ b/packages/core/src/amazonqGumby/chat/controller/messenger/messenger.ts @@ -11,7 +11,7 @@ import vscode from 'vscode' import { AuthFollowUpType, AuthMessageDataMap } from '../../../../amazonq/auth/model' import { JDKVersion, TransformationCandidateProject, transformByQState } from '../../../../codewhisperer/models/model' -import { FeatureAuthState } from '../../../../codewhisperer/util/authUtil' +import { AuthState } from '../../../../auth/auth2' import * as CodeWhispererConstants from '../../../../codewhisperer/models/constants' import { AppToWebViewMessageDispatcher, @@ -92,19 +92,15 @@ export class Messenger { this.dispatcher.sendUpdatePlaceholder(new UpdatePlaceholderMessage(tabID, newPlaceholder)) } - public async sendAuthNeededExceptionMessage(credentialState: FeatureAuthState, tabID: string) { + public async sendAuthNeededExceptionMessage(credentialState: AuthState, tabID: string) { let authType: AuthFollowUpType = 'full-auth' let message = AuthMessageDataMap[authType].message - switch (credentialState.amazonQ) { - case 'disconnected': + switch (credentialState) { + case 'notConnected': authType = 'full-auth' message = AuthMessageDataMap[authType].message break - case 'unsupported': - authType = 'use-supported-auth' - message = AuthMessageDataMap[authType].message - break case 'expired': authType = 're-auth' message = AuthMessageDataMap[authType].message diff --git a/packages/core/src/amazonqTest/app.ts b/packages/core/src/amazonqTest/app.ts index 6c638c13b71..1988ceb1d17 100644 --- a/packages/core/src/amazonqTest/app.ts +++ b/packages/core/src/amazonqTest/app.ts @@ -50,7 +50,7 @@ export function init(appContext: AmazonQAppInitContext) { appContext.registerWebViewToAppMessagePublisher(new MessagePublisher(testChatUIInputEventEmitter), 'testgen') const debouncedEvent = debounce(async () => { - const authenticated = (await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected' + const authenticated = AuthUtil.instance.getAuthState() === 'connected' let authenticatingSessionID = '' if (authenticated) { @@ -65,7 +65,7 @@ export function init(appContext: AmazonQAppInitContext) { messenger.sendAuthenticationUpdate(authenticated, [authenticatingSessionID]) }, 500) - AuthUtil.instance.secondaryAuth.onDidChangeActiveConnection(() => { + AuthUtil.instance.onDidChangeConnectionState(() => { return debouncedEvent() }) AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(() => { diff --git a/packages/core/src/amazonqTest/chat/controller/controller.ts b/packages/core/src/amazonqTest/chat/controller/controller.ts index 747cca57e8e..286ce8644ff 100644 --- a/packages/core/src/amazonqTest/chat/controller/controller.ts +++ b/packages/core/src/amazonqTest/chat/controller/controller.ts @@ -221,8 +221,8 @@ export class TestController { try { logger.debug(`Q - Test: Session created with id: ${session.tabID}`) - const authState = await AuthUtil.instance.getChatAuthState() - if (authState.amazonQ !== 'connected') { + const authState = AuthUtil.instance.getAuthState() + if (authState !== 'connected') { void this.messenger.sendAuthNeededExceptionMessage(authState, tabID) session.isAuthenticating = true return @@ -239,7 +239,7 @@ export class TestController { telemetry.amazonq_feedback.emit({ featureId: 'amazonQTest', amazonqConversationId: session.startTestGenerationRequestId, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, interactionType: message.vote, }) } @@ -487,8 +487,8 @@ export class TestController { userPrompt = message.prompt.slice(0, maxUserPromptLength) // check that the session is authenticated - const authState = await AuthUtil.instance.getChatAuthState() - if (authState.amazonQ !== 'connected') { + const authState = AuthUtil.instance.getAuthState() + if (authState !== 'connected') { void this.messenger.sendAuthNeededExceptionMessage(authState, tabID) session.isAuthenticating = true return diff --git a/packages/core/src/amazonqTest/chat/controller/messenger/messenger.ts b/packages/core/src/amazonqTest/chat/controller/messenger/messenger.ts index 5541ef389c5..17e88c2621b 100644 --- a/packages/core/src/amazonqTest/chat/controller/messenger/messenger.ts +++ b/packages/core/src/amazonqTest/chat/controller/messenger/messenger.ts @@ -9,7 +9,6 @@ */ import { AuthFollowUpType, AuthMessageDataMap } from '../../../../amazonq/auth/model' -import { FeatureAuthState } from '../../../../codewhisperer/util/authUtil' import { AppToWebViewMessageDispatcher, AuthNeededException, @@ -39,6 +38,7 @@ import { keys } from '../../../../shared/utilities/tsUtils' import { cancellingProgressField, testGenCompletedField } from '../../../models/constants' import { testGenState } from '../../../../codewhisperer/models/model' import { TelemetryHelper } from '../../../../codewhisperer/util/telemetryHelper' +import { AuthState } from '../../../../auth/auth2' export type UnrecoverableErrorType = 'no-project-found' | 'no-open-file-found' | 'invalid-file-type' @@ -122,19 +122,15 @@ export class Messenger { this.dispatcher.sendUpdatePromptProgress(new UpdatePromptProgressMessage(tabID, progressField)) } - public async sendAuthNeededExceptionMessage(credentialState: FeatureAuthState, tabID: string) { + public async sendAuthNeededExceptionMessage(credentialState: AuthState, tabID: string) { let authType: AuthFollowUpType = 'full-auth' let message = AuthMessageDataMap[authType].message - switch (credentialState.amazonQ) { - case 'disconnected': + switch (credentialState) { + case 'notConnected': authType = 'full-auth' message = AuthMessageDataMap[authType].message break - case 'unsupported': - authType = 'use-supported-auth' - message = AuthMessageDataMap[authType].message - break case 'expired': authType = 're-auth' message = AuthMessageDataMap[authType].message diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts new file mode 100644 index 00000000000..273a644ebbd --- /dev/null +++ b/packages/core/src/auth/auth2.ts @@ -0,0 +1,363 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as vscode from 'vscode' +import * as jose from 'jose' +import { + GetSsoTokenParams, + getSsoTokenRequestType, + GetSsoTokenResult, + IamIdentityCenterSsoTokenSource, + InvalidateSsoTokenParams, + invalidateSsoTokenRequestType, + ProfileKind, + UpdateProfileParams, + updateProfileRequestType, + SsoTokenChangedParams, + ssoTokenChangedRequestType, + AwsBuilderIdSsoTokenSource, + UpdateCredentialsParams, + AwsErrorCodes, + SsoTokenSourceKind, + listProfilesRequestType, + ListProfilesResult, + UpdateProfileResult, + InvalidateSsoTokenResult, + AuthorizationFlowKind, + CancellationToken, + CancellationTokenSource, + bearerCredentialsDeleteNotificationType, + bearerCredentialsUpdateRequestType, + SsoTokenChangedKind, + RequestType, + ResponseMessage, + NotificationType, + ConnectionMetadata, + getConnectionMetadataRequestType, +} from '@aws/language-server-runtimes/protocol' +import { LanguageClient } from 'vscode-languageclient' +import { getLogger } from '../shared/logger/logger' +import { ToolkitError } from '../shared/errors' +import { useDeviceFlow } from './sso/ssoAccessTokenProvider' +import { getCacheDir, getCacheFileWatcher, getFlareCacheFileName } from './sso/cache' +import { VSCODE_EXTENSION_ID } from '../shared/extensions' + +export const notificationTypes = { + updateBearerToken: new RequestType( + bearerCredentialsUpdateRequestType.method + ), + deleteBearerToken: new NotificationType(bearerCredentialsDeleteNotificationType.method), + getConnectionMetadata: new RequestType( + getConnectionMetadataRequestType.method + ), +} + +export type AuthState = 'notConnected' | 'connected' | 'expired' + +export type AuthStateEvent = { id: string; state: AuthState | 'refreshed' } + +export const LoginTypes = { + SSO: 'sso', + IAM: 'iam', +} as const +export type LoginType = (typeof LoginTypes)[keyof typeof LoginTypes] + +interface BaseLogin { + readonly loginType: LoginType +} + +export type cacheChangedEvent = 'delete' | 'create' + +export type Login = SsoLogin // TODO: add IamLogin type when supported + +export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoTokenSource + +/** + * Handles auth requests to the Identity Server in the Amazon Q LSP. + */ +export class LanguageClientAuth { + readonly #ssoCacheWatcher = getCacheFileWatcher(getCacheDir(), getFlareCacheFileName(VSCODE_EXTENSION_ID.amazonq)) + + constructor( + private readonly client: LanguageClient, + private readonly clientName: string, + public readonly encryptionKey: Buffer + ) {} + + public get cacheWatcher() { + return this.#ssoCacheWatcher + } + + getSsoToken( + tokenSource: TokenSource, + login: boolean = false, + cancellationToken?: CancellationToken + ): Promise { + return this.client.sendRequest( + getSsoTokenRequestType.method, + { + clientName: this.clientName, + source: tokenSource, + options: { + loginOnInvalidToken: login, + authorizationFlow: useDeviceFlow() ? AuthorizationFlowKind.DeviceCode : AuthorizationFlowKind.Pkce, + }, + } satisfies GetSsoTokenParams, + cancellationToken + ) + } + + updateProfile( + profileName: string, + startUrl: string, + region: string, + scopes: string[] + ): Promise { + return this.client.sendRequest(updateProfileRequestType.method, { + profile: { + kinds: [ProfileKind.SsoTokenProfile], + name: profileName, + settings: { + region, + sso_session: profileName, + }, + }, + ssoSession: { + name: profileName, + settings: { + sso_region: region, + sso_start_url: startUrl, + sso_registration_scopes: scopes, + }, + }, + } satisfies UpdateProfileParams) + } + + listProfiles() { + return this.client.sendRequest(listProfilesRequestType.method, {}) as Promise + } + + /** + * Returns a profile by name along with its linked sso_session. + * Does not currently exist as an API in the Identity Service. + */ + async getProfile(profileName: string) { + const response = await this.listProfiles() + const profile = response.profiles.find((profile) => profile.name === profileName) + const ssoSession = profile?.settings?.sso_session + ? response.ssoSessions.find((session) => session.name === profile!.settings!.sso_session) + : undefined + + return { profile, ssoSession } + } + + updateBearerToken(request: UpdateCredentialsParams) { + return this.client.sendRequest(bearerCredentialsUpdateRequestType.method, request) + } + + deleteBearerToken() { + return this.client.sendNotification(bearerCredentialsDeleteNotificationType.method) + } + + invalidateSsoToken(tokenId: string) { + return this.client.sendRequest(invalidateSsoTokenRequestType.method, { + ssoTokenId: tokenId, + } satisfies InvalidateSsoTokenParams) as Promise + } + + registerSsoTokenChangedHandler(ssoTokenChangedHandler: (params: SsoTokenChangedParams) => any) { + this.client.onNotification(ssoTokenChangedRequestType.method, ssoTokenChangedHandler) + } + + registerCacheWatcher(cacheChangedHandler: (event: cacheChangedEvent) => any) { + this.cacheWatcher.onDidCreate(() => cacheChangedHandler('create')) + this.cacheWatcher.onDidDelete(() => cacheChangedHandler('delete')) + } +} + +/** + * Manages an SSO connection. + */ +export class SsoLogin implements BaseLogin { + readonly loginType = LoginTypes.SSO + private readonly eventEmitter = new vscode.EventEmitter() + + // Cached information from the identity server for easy reference + private ssoTokenId: string | undefined + private connectionState: AuthState = 'notConnected' + private _data: { startUrl: string; region: string } | undefined + + private cancellationToken: CancellationTokenSource | undefined + + constructor( + public readonly profileName: string, + private readonly lspAuth: LanguageClientAuth + ) { + lspAuth.registerSsoTokenChangedHandler((params: SsoTokenChangedParams) => this.ssoTokenChangedHandler(params)) + } + + get data() { + return this._data + } + + async login(opts: { startUrl: string; region: string; scopes: string[] }) { + await this.updateProfile(opts) + return this._getSsoToken(true) + } + + async reauthenticate() { + if (this.connectionState === 'notConnected') { + throw new ToolkitError('Cannot reauthenticate when not connected.') + } + return this._getSsoToken(true) + } + + async logout() { + if (this.ssoTokenId) { + await this.lspAuth.invalidateSsoToken(this.ssoTokenId) + } + this.updateConnectionState('notConnected') + this._data = undefined + // TODO: DeleteProfile api in Identity Service (this doesn't exist yet) + } + + async getProfile() { + return await this.lspAuth.getProfile(this.profileName) + } + + async updateProfile(opts: { startUrl: string; region: string; scopes: string[] }) { + await this.lspAuth.updateProfile(this.profileName, opts.startUrl, opts.region, opts.scopes) + this._data = { + startUrl: opts.startUrl, + region: opts.region, + } + } + + /** + * Restore the connection state and connection details to memory, if they exist. + */ + async restore() { + const sessionData = await this.getProfile() + const ssoSession = sessionData?.ssoSession?.settings + if (ssoSession?.sso_region && ssoSession?.sso_start_url) { + this._data = { + startUrl: ssoSession.sso_start_url, + region: ssoSession.sso_region, + } + } + + try { + await this._getSsoToken(false) + } catch (err) { + getLogger().error('Restoring connection failed: %s', err) + } + } + + /** + * Cancels running active login flows. + */ + cancelLogin() { + this.cancellationToken?.cancel() + this.cancellationToken?.dispose() + this.cancellationToken = undefined + } + + /** + * Returns both the decrypted access token and the payload to send to the `updateCredentials` LSP API + * with encrypted token + */ + async getToken() { + const response = await this._getSsoToken(false) + const decryptedKey = await jose.compactDecrypt(response.ssoToken.accessToken, this.lspAuth.encryptionKey) + return { + token: decryptedKey.plaintext.toString().replaceAll('"', ''), + updateCredentialsParams: response.updateCredentialsParams, + } + } + + /** + * Returns the response from `getSsoToken` LSP API and sets the connection state based on the errors/result + * of the call. + */ + private async _getSsoToken(login: boolean) { + let response: GetSsoTokenResult + this.cancellationToken = new CancellationTokenSource() + + try { + response = await this.lspAuth.getSsoToken( + { + /** + * Note that we do not use SsoTokenSourceKind.AwsBuilderId here. + * This is because it does not leave any state behind on disk, so + * we cannot infer that a builder ID connection exists via the + * Identity Server alone. + */ + kind: SsoTokenSourceKind.IamIdentityCenter, + profileName: this.profileName, + } satisfies IamIdentityCenterSsoTokenSource, + login, + this.cancellationToken.token + ) + } catch (err: any) { + switch (err.data?.awsErrorCode) { + case AwsErrorCodes.E_CANCELLED: + case AwsErrorCodes.E_SSO_SESSION_NOT_FOUND: + case AwsErrorCodes.E_PROFILE_NOT_FOUND: + case AwsErrorCodes.E_INVALID_SSO_TOKEN: + this.updateConnectionState('notConnected') + break + case AwsErrorCodes.E_CANNOT_REFRESH_SSO_TOKEN: + this.updateConnectionState('expired') + break + // TODO: implement when identity server emits E_NETWORK_ERROR, E_FILESYSTEM_ERROR + // case AwsErrorCodes.E_NETWORK_ERROR: + // case AwsErrorCodes.E_FILESYSTEM_ERROR: + // // do stuff, probably nothing at all + // break + default: + getLogger().error('SsoLogin: unknown error when requesting token: %s', err) + break + } + throw err + } finally { + this.cancellationToken?.dispose() + this.cancellationToken = undefined + } + + this.ssoTokenId = response.ssoToken.id + this.updateConnectionState('connected') + return response + } + + getConnectionState() { + return this.connectionState + } + + onDidChangeConnectionState(handler: (e: AuthStateEvent) => any) { + return this.eventEmitter.event(handler) + } + + private updateConnectionState(state: AuthState) { + const oldState = this.connectionState + const newState = state + + this.connectionState = newState + + if (oldState !== newState) { + this.eventEmitter.fire({ id: this.profileName, state: this.connectionState }) + } + } + + private ssoTokenChangedHandler(params: SsoTokenChangedParams) { + if (params.ssoTokenId === this.ssoTokenId) { + if (params.kind === SsoTokenChangedKind.Expired) { + this.updateConnectionState('expired') + return + } else if (params.kind === SsoTokenChangedKind.Refreshed) { + this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' }) + } + } + } +} diff --git a/packages/core/src/auth/index.ts b/packages/core/src/auth/index.ts index 02a0067be45..87a6877844e 100644 --- a/packages/core/src/auth/index.ts +++ b/packages/core/src/auth/index.ts @@ -23,4 +23,8 @@ export { export { Auth } from './auth' export { CredentialsStore } from './credentials/store' export { LoginManager } from './deprecated/loginManager' -export * as AuthUtils from './utils' +export * as constants from './sso/constants' +export * as cache from './sso/cache' +export * as authUtils from './utils' +export * as auth2 from './auth2' +export * as SsoAccessTokenProvider from './sso/ssoAccessTokenProvider' diff --git a/packages/core/src/auth/sso/cache.ts b/packages/core/src/auth/sso/cache.ts index 7d43c07da35..f9d62c50305 100644 --- a/packages/core/src/auth/sso/cache.ts +++ b/packages/core/src/auth/sso/cache.ts @@ -45,8 +45,8 @@ export function getCache(directory = getCacheDir()): SsoCache { } } -export function getCacheFileWatcher(directory = getCacheDir()) { - const watcher = vscode.workspace.createFileSystemWatcher(new vscode.RelativePattern(directory, '*.json')) +export function getCacheFileWatcher(directory = getCacheDir(), file = '*.json') { + const watcher = vscode.workspace.createFileSystemWatcher(new vscode.RelativePattern(directory, file)) globals.context.subscriptions.push(watcher) return watcher } @@ -126,7 +126,7 @@ export function getTokenCache(directory = getCacheDir()): KeyedCache return mapCache(cache, read, write) } -function getTokenCacheFile(ssoCacheDir: string, key: string): string { +export function getTokenCacheFile(ssoCacheDir: string, key: string): string { const encoded = encodeURI(key) // Per the spec: 'SSO Login Token Flow' the access token must be // cached as the SHA1 hash of the bytes of the UTF-8 encoded @@ -145,7 +145,7 @@ function getTokenCacheFile(ssoCacheDir: string, key: string): string { return path.join(ssoCacheDir, `${hashedKey}.json`) } -function getRegistrationCacheFile(ssoCacheDir: string, key: RegistrationKey): string { +export function getRegistrationCacheFile(ssoCacheDir: string, key: RegistrationKey): string { const hash = (startUrl: string, scopes: string[]) => { const shasum = crypto.createHash('sha256') shasum.update(startUrl) @@ -158,3 +158,19 @@ function getRegistrationCacheFile(ssoCacheDir: string, key: RegistrationKey): st const suffix = `${key.region}${key.scopes && key.scopes.length > 0 ? `-${hash(key.startUrl, key.scopes)}` : ''}` return path.join(ssoCacheDir, `aws-toolkit-vscode-client-id-${suffix}.json`) } + +/** + * Returns the cache file name that Flare identity server uses for SSO token and registration + * + * @param key - The key to use for the new registration cache file. + * See https://github.com/aws/language-servers/blob/c10819ea2c25ce564c75fb43a6792f3c919b757a/server/aws-lsp-identity/src/sso/cache/fileSystemSsoCache.ts + * @returns File name of the Flare cache file + */ +export function getFlareCacheFileName(key: string) { + const hash = (str: string) => { + const hasher = crypto.createHash('sha1') + return hasher.update(str).digest('hex') + } + + return hash(key) + '.json' +} diff --git a/packages/core/src/auth/sso/constants.ts b/packages/core/src/auth/sso/constants.ts index 0e6bb082d7e..14d2382a692 100644 --- a/packages/core/src/auth/sso/constants.ts +++ b/packages/core/src/auth/sso/constants.ts @@ -10,6 +10,7 @@ export const builderIdStartUrl = 'https://view.awsapps.com/start' export const internalStartUrl = 'https://amzn.awsapps.com/start' +export const builderIdRegion = 'us-east-1' /** * Doc: https://docs.aws.amazon.com/singlesignon/latest/userguide/howtochangeURL.html diff --git a/packages/core/src/auth/sso/ssoAccessTokenProvider.ts b/packages/core/src/auth/sso/ssoAccessTokenProvider.ts index e753fb2ef90..7caf2638e1b 100644 --- a/packages/core/src/auth/sso/ssoAccessTokenProvider.ts +++ b/packages/core/src/auth/sso/ssoAccessTokenProvider.ts @@ -289,17 +289,7 @@ export abstract class SsoAccessTokenProvider { profile: Pick, cache = getCache(), oidc: OidcClient = OidcClient.create(profile.region), - reAuthState?: ReAuthState, - useDeviceFlow: () => boolean = () => { - /** - * Device code flow is neccessary when: - * 1. We are in a workspace connected through ssh (codecatalyst, etc) - * 2. We are connected to a remote backend through the web browser (code server, openshift dev spaces) - * - * Since we are unable to serve the final authorization page - */ - return getExtRuntimeContext().extensionHost === 'remote' - } + reAuthState?: ReAuthState ) { if (DevSettings.instance.get('webAuth', false) && getExtRuntimeContext().extensionHost === 'webworker') { return new WebAuthorization(profile, cache, oidc, reAuthState) @@ -400,6 +390,17 @@ function getSessionDuration(id: string) { return creationDate !== undefined ? globals.clock.Date.now() - creationDate : undefined } +export function useDeviceFlow(): boolean { + /** + * Device code flow is neccessary when: + * 1. We are in a workspace connected through ssh (codecatalyst, etc) + * 2. We are connected to a remote backend through the web browser (code server, openshift dev spaces) + * + * Since we are unable to serve the final authorization page + */ + return getExtRuntimeContext().extensionHost === 'remote' +} + /** * SSO "device code" flow (RFC: https://tools.ietf.org/html/rfc8628) * 1. Get a client id (SSO-OIDC identifier, formatted per RFC6749). diff --git a/packages/core/src/auth/utils.ts b/packages/core/src/auth/utils.ts index 9da35fa06e5..1592774a338 100644 --- a/packages/core/src/auth/utils.ts +++ b/packages/core/src/auth/utils.ts @@ -45,8 +45,8 @@ import { Commands, placeholder } from '../shared/vscode/commands2' import { Auth } from './auth' import { validateIsNewSsoUrl, validateSsoUrlFormat } from './sso/validation' import { getLogger } from '../shared/logger/logger' -import { AuthUtil, isValidAmazonQConnection, isValidCodeWhispererCoreConnection } from '../codewhisperer/util/authUtil' import { AuthFormId } from '../login/webview/vue/types' +import { amazonQScopes, AuthUtil } from '../codewhisperer/util/authUtil' import { extensionVersion } from '../shared/vscode/env' import { CommonAuthWebview } from '../login/webview/vue/backend' import { AuthSource } from '../login/webview/util' @@ -585,7 +585,7 @@ export async function hasIamCredentials( return (await allConnections()).some(isIamConnection) } -export type SsoKind = 'any' | 'codewhisperer' | 'codecatalyst' +export type SsoKind = 'any' | 'codecatalyst' /** * Returns true if an Identity Center SSO connection exists. @@ -606,11 +606,6 @@ export async function findSsoConnections( ): Promise { let predicate: (c?: Connection) => boolean switch (kind) { - case 'codewhisperer': - predicate = (conn?: Connection) => { - return isIdcSsoConnection(conn) && isValidCodeWhispererCoreConnection(conn) - } - break case 'codecatalyst': predicate = (conn?: Connection) => { return isIdcSsoConnection(conn) && isValidCodeCatalystConnection(conn) @@ -622,7 +617,7 @@ export async function findSsoConnections( return (await allConnections()).filter(predicate).filter(isIdcSsoConnection) } -export type BuilderIdKind = 'any' | 'codewhisperer' | 'codecatalyst' +export type BuilderIdKind = 'any' | 'codecatalyst' /** * Returns true if a Builder ID connection exists. @@ -643,11 +638,6 @@ async function findBuilderIdConnections( ): Promise { let predicate: (c?: Connection) => boolean switch (kind) { - case 'codewhisperer': - predicate = (conn?: Connection) => { - return isBuilderIdConnection(conn) && isValidCodeWhispererCoreConnection(conn) - } - break case 'codecatalyst': predicate = (conn?: Connection) => { return isBuilderIdConnection(conn) && isValidCodeCatalystConnection(conn) @@ -656,7 +646,7 @@ async function findBuilderIdConnections( case 'any': predicate = isBuilderIdConnection } - return (await allConnections()).filter(predicate).filter(isAnySsoConnection) + return (await allConnections()).filter(predicate!).filter(isAnySsoConnection) } /** @@ -803,7 +793,7 @@ export function getAuthFormIdsFromConnection(conn?: Connection): AuthFormId[] { if (isValidCodeCatalystConnection(conn)) { authIds.push(`${connType}CodeCatalyst`) } - if (isValidAmazonQConnection(conn)) { + if (hasScopes(conn, amazonQScopes)) { authIds.push(`${connType}CodeWhisperer`) } @@ -818,9 +808,9 @@ export function initializeCredentialsProviderManager() { export async function getAuthType() { let authType: CredentialSourceId | undefined = undefined - if (AuthUtil.instance.isEnterpriseSsoInUse() && AuthUtil.instance.isConnectionValid()) { + if (AuthUtil.instance.isIdcConnection()) { authType = 'iamIdentityCenter' - } else if (AuthUtil.instance.isBuilderIdInUse() && AuthUtil.instance.isConnectionValid()) { + } else if (AuthUtil.instance.isBuilderIdConnection()) { authType = 'awsId' } return authType diff --git a/packages/core/src/codewhisperer/activation.ts b/packages/core/src/codewhisperer/activation.ts index e52e08bb98b..ab260dfc5e1 100644 --- a/packages/core/src/codewhisperer/activation.ts +++ b/packages/core/src/codewhisperer/activation.ts @@ -5,8 +5,6 @@ import * as vscode from 'vscode' import * as nls from 'vscode-nls' -import { getTabSizeSetting } from '../shared/utilities/editorUtilities' -import * as EditorContext from './util/editorContext' import * as CodeWhispererConstants from './models/constants' import { CodeSuggestionsState, @@ -16,7 +14,6 @@ import { CodeScanIssue, CodeIssueGroupingStrategyState, } from './models/model' -import { acceptSuggestion } from './commands/onInlineAcceptance' import { CodeWhispererSettings } from './util/codewhispererSettings' import { ExtContext } from '../shared/extensions' import { CodeWhispererTracker } from './tracker/codewhispererTracker' @@ -30,7 +27,6 @@ import { showLearnMore, showSsoSignIn, showFreeTierLimit, - updateReferenceLog, showIntroduction, reconnect, openSecurityIssuePanel, @@ -64,20 +60,15 @@ import { updateSecurityDiagnosticCollection, } from './service/diagnosticsProvider' import { SecurityPanelViewProvider, openEditorAtRange } from './views/securityPanelViewProvider' -import { RecommendationHandler } from './service/recommendationHandler' import { Commands, registerCommandErrorHandler, registerDeclaredCommands } from '../shared/vscode/commands2' -import { InlineCompletionService, refreshStatusBar } from './service/inlineCompletionService' -import { isInlineCompletionEnabled } from './util/commonUtil' import { AuthUtil } from './util/authUtil' import { ImportAdderProvider } from './service/importAdderProvider' -import { TelemetryHelper } from './util/telemetryHelper' import { openUrl } from '../shared/utilities/vsCodeUtils' -import { notifyNewCustomizations, onProfileChangedListener } from './util/customizationUtil' +import { onProfileChangedListener } from './util/customizationUtil' import { CodeWhispererCommandBackend, CodeWhispererCommandDeclarations } from './commands/gettingStartedPageCommands' import { SecurityIssueHoverProvider } from './service/securityIssueHoverProvider' import { SecurityIssueCodeActionProvider } from './service/securityIssueCodeActionProvider' import { listCodeWhispererCommands } from './ui/statusBarMenu' -import { Container } from './service/serviceContainer' import { debounceStartSecurityScan } from './commands/startSecurityScan' import { securityScanLanguageContext } from './util/securityScanLanguageContext' import { registerWebviewErrorHandler } from '../webviews/server' @@ -91,7 +82,6 @@ import { setContext } from '../shared/vscode/setContext' import { syncSecurityIssueWebview } from './views/securityIssue/securityIssueWebview' import { detectCommentAboveLine } from '../shared/utilities/commentUtils' import { activateEditTracking } from './nextEditPrediction/activation' -import { notifySelectDeveloperProfile } from './region/utils' let localize: nls.LocalizeFunc @@ -101,10 +91,6 @@ export async function activate(context: ExtContext): Promise { // Import old CodeWhisperer settings into Amazon Q await CodeWhispererSettings.instance.importSettings() - // initialize AuthUtil earlier to make sure it can listen to connection change events. - const auth = AuthUtil.instance - auth.initCodeWhispererHooks() - // TODO: is this indirection useful? registerDeclaredCommands( context.extensionContext.subscriptions, @@ -137,25 +123,20 @@ export async function activate(context: ExtContext): Promise { const client = new codewhispererClient.DefaultCodeWhispererClient() // Service initialization - const container = Container.instance ReferenceInlineProvider.instance ImportAdderProvider.instance context.extensionContext.subscriptions.push( // register toolkit api callback registerToolkitApiCallback.register(), - signoutCodeWhisperer.register(auth), + signoutCodeWhisperer.register(), /** * Configuration change */ vscode.workspace.onDidChangeConfiguration(async (configurationChangeEvent) => { - if (configurationChangeEvent.affectsConfiguration('editor.tabSize')) { - EditorContext.updateTabSize(getTabSizeSetting()) - } - if (configurationChangeEvent.affectsConfiguration('amazonQ.showCodeWithReferences')) { ReferenceLogViewProvider.instance.update() - if (auth.isEnterpriseSsoInUse()) { + if (AuthUtil.instance.isIdcConnection()) { await vscode.window .showInformationMessage( CodeWhispererConstants.ssoConfigAlertMessage, @@ -170,7 +151,7 @@ export async function activate(context: ExtContext): Promise { } if (configurationChangeEvent.affectsConfiguration('amazonQ.shareContentWithAWS')) { - if (auth.isEnterpriseSsoInUse()) { + if (AuthUtil.instance.isIdcConnection()) { await vscode.window .showInformationMessage( CodeWhispererConstants.ssoConfigAlertMessageShareData, @@ -215,20 +196,21 @@ export async function activate(context: ExtContext): Promise { await openSettings('amazonQ') } }), - Commands.register('aws.amazonq.refreshAnnotation', async (forceProceed: boolean) => { - telemetry.record({ - traceId: TelemetryHelper.instance.traceId, - }) - - const editor = vscode.window.activeTextEditor - if (editor) { - if (forceProceed) { - await container.lineAnnotationController.refresh(editor, 'codewhisperer', true) - } else { - await container.lineAnnotationController.refresh(editor, 'codewhisperer') - } - } - }), + // TODO port this to lsp + // Commands.register('aws.amazonq.refreshAnnotation', async (forceProceed: boolean) => { + // telemetry.record({ + // traceId: TelemetryHelper.instance.traceId, + // }) + + // const editor = vscode.window.activeTextEditor + // if (editor) { + // if (forceProceed) { + // await container.lineAnnotationController.refresh(editor, 'codewhisperer', true) + // } else { + // await container.lineAnnotationController.refresh(editor, 'codewhisperer') + // } + // } + // }), // show introduction showIntroduction.register(), // toggle code suggestions @@ -251,10 +233,6 @@ export async function activate(context: ExtContext): Promise { showLearnMore.register(), // show free tier limit showFreeTierLimit.register(), - // update reference log instance - updateReferenceLog.register(), - // refresh codewhisperer status bar - refreshStatusBar.register(), // generate code fix generateFix.register(client, context), // regenerate code fix @@ -300,22 +278,10 @@ export async function activate(context: ExtContext): Promise { // notify new customizations notifyNewCustomizationsCmd.register(), selectRegionProfileCommand.register(), - /** - * On recommendation acceptance - */ - acceptSuggestion.register(context), // direct CodeWhisperer connection setup with customization connectWithCustomization.register(), - // on text document close. - vscode.workspace.onDidCloseTextDocument((e) => { - if (isInlineCompletionEnabled() && e.uri.fsPath !== InlineCompletionService.instance.filePath()) { - return - } - RecommendationHandler.instance.reportUserDecisions(-1) - }), - vscode.languages.registerHoverProvider( [...CodeWhispererConstants.platformLanguageIds], ReferenceHoverProvider.instance @@ -340,36 +306,26 @@ export async function activate(context: ExtContext): Promise { SecurityIssueCodeActionProvider.instance ), vscode.commands.registerCommand('aws.amazonq.openEditorAtRange', openEditorAtRange), - auth.regionProfileManager.onDidChangeRegionProfile(onProfileChangedListener) + AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(onProfileChangedListener) ) // run the auth startup code with context for telemetry await telemetry.function_call.run( async () => { - await auth.restore() - await auth.clearExtraConnections() - - if (auth.isConnectionExpired()) { - auth.showReauthenticatePrompt().catch((e) => { + if (AuthUtil.instance.isConnectionExpired()) { + AuthUtil.instance.showReauthenticatePrompt().catch((e) => { const defaulMsg = localize('AWS.generic.message.error', 'Failed to reauth:') void logAndShowError(localize, e, 'showReauthenticatePrompt', defaulMsg) }) - if (auth.isEnterpriseSsoInUse()) { - await auth.notifySessionConfiguration() + if (AuthUtil.instance.isIdcConnection()) { + await AuthUtil.instance.notifySessionConfiguration() } } - - if (auth.requireProfileSelection()) { - await notifySelectDeveloperProfile() - } }, { emit: false, functionId: { name: 'activateCwCore' } } ) - if (auth.isValidEnterpriseSsoInUse()) { - await notifyNewCustomizations() - } - if (auth.isBuilderIdInUse()) { + if (AuthUtil.instance.isBuilderIdConnection()) { await CodeScansState.instance.setScansEnabled(false) } @@ -384,8 +340,8 @@ export async function activate(context: ExtContext): Promise { return ( (isScansEnabled ?? CodeScansState.instance.isScansEnabled()) && !CodeScansState.instance.isMonthlyQuotaExceeded() && - auth.isConnectionValid() && - !auth.isBuilderIdInUse() && + AuthUtil.instance.isConnected() && + !AuthUtil.instance.isBuilderIdConnection() && editor && editor.document.uri.scheme === 'file' && securityScanLanguageContext.isLanguageSupported(editor.document.languageId) @@ -473,7 +429,6 @@ export async function activate(context: ExtContext): Promise { }) await Commands.tryExecute('aws.amazonq.refreshConnectionCallback') - container.ready() function setSubscriptionsForCodeIssues() { context.extensionContext.subscriptions.push( @@ -511,8 +466,8 @@ export async function activate(context: ExtContext): Promise { } export async function shutdown() { - RecommendationHandler.instance.reportUserDecisions(-1) await CodeWhispererTracker.getTracker().shutdown() + AuthUtil.instance.regionProfileManager.globalStatePoller.kill() } function toggleIssuesVisibility(visibleCondition: (issue: CodeScanIssue, filePath: string) => boolean) { diff --git a/packages/core/src/codewhisperer/client/codewhisperer.ts b/packages/core/src/codewhisperer/client/codewhisperer.ts index 35f699b24c2..0a473dfdccd 100644 --- a/packages/core/src/codewhisperer/client/codewhisperer.ts +++ b/packages/core/src/codewhisperer/client/codewhisperer.ts @@ -13,7 +13,6 @@ import { hasVendedIamCredentials } from '../../auth/auth' import { CodeWhispererSettings } from '../util/codewhispererSettings' import { PromiseResult } from 'aws-sdk/lib/request' import { AuthUtil } from '../util/authUtil' -import { isSsoConnection } from '../../auth/connection' import apiConfig = require('./service-2.json') import userApiConfig = require('./user-service-2.json') import { session } from '../util/codeWhispererSession' @@ -84,6 +83,8 @@ export type Imports = CodeWhispererUserClient.Imports export class DefaultCodeWhispererClient { private async createSdkClient(): Promise { + throw new Error('Do not call this function until IAM is supported by LSP identity server') + const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled() const cwsprConfig = getCodewhispererConfig() return (await globals.sdkClientBuilder.createAwsService( @@ -91,7 +92,7 @@ export class DefaultCodeWhispererClient { { apiConfig: apiConfig, region: cwsprConfig.region, - credentials: await AuthUtil.instance.getCredentials(), + credentials: undefined, endpoint: cwsprConfig.endpoint, onRequestSetup: [ (req) => { @@ -126,7 +127,7 @@ export class DefaultCodeWhispererClient { async createUserSdkClient(maxRetries?: number): Promise { const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled() session.setFetchCredentialStart() - const bearerToken = await AuthUtil.instance.getBearerToken() + const bearerToken = await AuthUtil.instance.getToken() session.setSdkApiCallStart() const cwsprConfig = getCodewhispererConfig() return (await globals.sdkClientBuilder.createAwsService( @@ -156,7 +157,7 @@ export class DefaultCodeWhispererClient { } private isBearerTokenAuth(): boolean { - return isSsoConnection(AuthUtil.instance.conn) + return AuthUtil.instance.isConnected() // TODO: Handle IAM credentials } public async generateRecommendations( @@ -230,7 +231,7 @@ export class DefaultCodeWhispererClient { }, profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, } - if (!AuthUtil.instance.isValidEnterpriseSsoInUse() && !globals.telemetry.telemetryEnabled) { + if (!AuthUtil.instance.isIdcConnection() && !globals.telemetry.telemetryEnabled) { return } const response = await (await this.createUserSdkClient()).sendTelemetryEvent(requestWithCommonFields).promise() diff --git a/packages/core/src/codewhisperer/commands/basicCommands.ts b/packages/core/src/codewhisperer/commands/basicCommands.ts index 7fe6078a1d7..1ecdf641fb4 100644 --- a/packages/core/src/codewhisperer/commands/basicCommands.ts +++ b/packages/core/src/codewhisperer/commands/basicCommands.ts @@ -52,7 +52,6 @@ import { removeDiagnostic } from '../service/diagnosticsProvider' import { SsoAccessTokenProvider } from '../../auth/sso/ssoAccessTokenProvider' import { ToolkitError, getErrorMsg, getTelemetryReason, getTelemetryReasonDesc } from '../../shared/errors' import { isRemoteWorkspace } from '../../shared/vscode/env' -import { isBuilderIdConnection } from '../../auth/connection' import globals from '../../shared/extensionGlobals' import { getVscodeCliPath } from '../../shared/utilities/pathFind' import { tryRun } from '../../shared/utilities/pathFind' @@ -114,7 +113,7 @@ export const toggleCodeScans = Commands.declare( { id: 'aws.codeWhisperer.toggleCodeScan', compositeKey: { 1: 'source' } }, (scansState: CodeScansState) => async (_: VsCodeCommandArg, source: CodeWhispererSource) => { await telemetry.aws_modifySetting.run(async (span) => { - if (isBuilderIdConnection(AuthUtil.instance.conn)) { + if (AuthUtil.instance.isBuilderIdConnection()) { throw new Error(`Auto-scans are not supported with the Amazon Builder ID connection.`) } span.record({ @@ -145,7 +144,7 @@ export const showReferenceLog = Commands.declare( if (_ !== placeholder) { source = 'ellipsesMenu' } - await vscode.commands.executeCommand('workbench.view.extension.aws-codewhisperer-reference-log') + await vscode.commands.executeCommand(`${ReferenceLogViewProvider.viewType}.focus`) } ) @@ -239,7 +238,7 @@ export const showFileScan = Commands.declare( export const selectCustomizationPrompt = Commands.declare( { id: 'aws.amazonq.selectCustomization', compositeKey: { 1: 'source' } }, () => async (_: VsCodeCommandArg, source: CodeWhispererSource) => { - if (isBuilderIdConnection(AuthUtil.instance.conn)) { + if (AuthUtil.instance.isBuilderIdConnection()) { throw new Error(`Select Customizations are not supported with the Amazon Builder ID connection.`) } telemetry.ui_click.emit({ elementId: 'cw_selectCustomization_Cta' }) @@ -376,7 +375,7 @@ export const openSecurityIssuePanel = Commands.declare( findingId: targetIssue.findingId, detectorId: targetIssue.detectorId, ruleId: targetIssue.ruleId, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, autoDetected: targetIssue.autoDetected, }) TelemetryHelper.instance.sendCodeScanRemediationsEvent( @@ -465,7 +464,7 @@ export const applySecurityFix = Commands.declare( ruleId: targetIssue.ruleId, component: targetSource, result: 'Succeeded', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, codeFixAction: 'applyFix', autoDetected: targetIssue.autoDetected, codewhispererCodeScanJobId: targetIssue.scanJobId, @@ -594,8 +593,8 @@ export const applySecurityFix = Commands.declare( export const signoutCodeWhisperer = Commands.declare( { id: 'aws.amazonq.signout', compositeKey: { 1: 'source' } }, - (auth: AuthUtil) => async (_: VsCodeCommandArg, source: CodeWhispererSource) => { - await auth.secondaryAuth.deleteConnection() + () => async (_: VsCodeCommandArg, source: CodeWhispererSource) => { + await AuthUtil.instance.logout() SecurityIssueTreeViewProvider.instance.refresh() return focusAmazonQPanel.execute(placeholder, source) } @@ -652,14 +651,13 @@ export const registerToolkitApiCallback = Commands.declare( if (_toolkitApi) { registerToolkitApiCallbackOnce() // Declare current conn immediately - const currentConn = AuthUtil.instance.conn - if (currentConn?.type === 'sso') { + if (AuthUtil.instance.isConnected() && AuthUtil.instance.isSsoSession()) { _toolkitApi.declareConnection( { - type: currentConn.type, - ssoRegion: currentConn.ssoRegion, - startUrl: currentConn.startUrl, - id: currentConn.id, + type: 'sso', + ssoRegion: AuthUtil.instance.connection?.region, + startUrl: AuthUtil.instance.connection?.startUrl, + id: AuthUtil.instance.profileName, } as AwsConnection, 'Amazon Q' ) @@ -855,7 +853,7 @@ export const ignoreAllIssues = Commands.declare( telemetry.record({ component: targetSource, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, detectorId: targetIssue.detectorId, findingId: targetIssue.findingId, ruleId: targetIssue.ruleId, @@ -890,7 +888,7 @@ export const ignoreIssue = Commands.declare( telemetry.record({ component: targetSource, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, detectorId: targetIssue.detectorId, findingId: targetIssue.findingId, ruleId: targetIssue.ruleId, diff --git a/packages/core/src/codewhisperer/commands/invokeRecommendation.ts b/packages/core/src/codewhisperer/commands/invokeRecommendation.ts deleted file mode 100644 index 37fcb965774..00000000000 --- a/packages/core/src/codewhisperer/commands/invokeRecommendation.ts +++ /dev/null @@ -1,45 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import { vsCodeState, ConfigurationEntry } from '../models/model' -import { resetIntelliSenseState } from '../util/globalStateUtil' -import { DefaultCodeWhispererClient } from '../client/codewhisperer' -import { RecommendationHandler } from '../service/recommendationHandler' -import { session } from '../util/codeWhispererSession' -import { RecommendationService } from '../service/recommendationService' - -/** - * This function is for manual trigger CodeWhisperer - */ - -export async function invokeRecommendation( - editor: vscode.TextEditor, - client: DefaultCodeWhispererClient, - config: ConfigurationEntry -) { - if (!editor || !config.isManualTriggerEnabled) { - return - } - - /** - * Skip when output channel gains focus and invoke - */ - if (editor.document.languageId === 'Log') { - return - } - /** - * When using intelliSense, if invocation position changed, reject previous active recommendations - */ - if (vsCodeState.isIntelliSenseActive && editor.selection.active !== session.startPos) { - resetIntelliSenseState( - config.isManualTriggerEnabled, - config.isAutomatedTriggerEnabled, - RecommendationHandler.instance.isValidResponse() - ) - } - - await RecommendationService.instance.generateRecommendation(client, editor, 'OnDemand', config, undefined) -} diff --git a/packages/core/src/codewhisperer/commands/onAcceptance.ts b/packages/core/src/codewhisperer/commands/onAcceptance.ts deleted file mode 100644 index e13c197cefd..00000000000 --- a/packages/core/src/codewhisperer/commands/onAcceptance.ts +++ /dev/null @@ -1,85 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import { vsCodeState, OnRecommendationAcceptanceEntry } from '../models/model' -import { runtimeLanguageContext } from '../util/runtimeLanguageContext' -import { CodeWhispererTracker } from '../tracker/codewhispererTracker' -import { CodeWhispererCodeCoverageTracker } from '../tracker/codewhispererCodeCoverageTracker' -import { getLogger } from '../../shared/logger/logger' -import { handleExtraBrackets } from '../util/closingBracketUtil' -import { RecommendationHandler } from '../service/recommendationHandler' -import { ReferenceLogViewProvider } from '../service/referenceLogViewProvider' -import { ReferenceHoverProvider } from '../service/referenceHoverProvider' -import path from 'path' - -/** - * This function is called when user accepts a intelliSense suggestion or an inline suggestion - */ -export async function onAcceptance(acceptanceEntry: OnRecommendationAcceptanceEntry) { - RecommendationHandler.instance.cancelPaginatedRequest() - /** - * Format document - */ - if (acceptanceEntry.editor) { - const languageContext = runtimeLanguageContext.getLanguageContext( - acceptanceEntry.editor.document.languageId, - path.extname(acceptanceEntry.editor.document.fileName) - ) - const start = acceptanceEntry.range.start - const end = acceptanceEntry.range.end - - // codewhisperer will be doing editing while formatting. - // formatting should not trigger consoals auto trigger - vsCodeState.isCodeWhispererEditing = true - /** - * Mitigation to right context handling mainly for auto closing bracket use case - */ - try { - await handleExtraBrackets(acceptanceEntry.editor, end, start) - } catch (error) { - getLogger().error(`${error} in handleAutoClosingBrackets`) - } - // move cursor to end of suggestion before doing code format - // after formatting, the end position will still be editor.selection.active - acceptanceEntry.editor.selection = new vscode.Selection(end, end) - - vsCodeState.isCodeWhispererEditing = false - CodeWhispererTracker.getTracker().enqueue({ - time: new Date(), - fileUrl: acceptanceEntry.editor.document.uri, - originalString: acceptanceEntry.editor.document.getText(new vscode.Range(start, end)), - startPosition: start, - endPosition: end, - requestId: acceptanceEntry.requestId, - sessionId: acceptanceEntry.sessionId, - index: acceptanceEntry.acceptIndex, - triggerType: acceptanceEntry.triggerType, - completionType: acceptanceEntry.completionType, - language: languageContext.language, - }) - const insertedCoderange = new vscode.Range(start, end) - CodeWhispererCodeCoverageTracker.getTracker(languageContext.language)?.countAcceptedTokens( - insertedCoderange, - acceptanceEntry.editor.document.getText(insertedCoderange), - acceptanceEntry.editor.document.fileName - ) - if (acceptanceEntry.references !== undefined) { - const referenceLog = ReferenceLogViewProvider.getReferenceLog( - acceptanceEntry.recommendation, - acceptanceEntry.references, - acceptanceEntry.editor - ) - ReferenceLogViewProvider.instance.addReferenceLog(referenceLog) - ReferenceHoverProvider.instance.addCodeReferences( - acceptanceEntry.recommendation, - acceptanceEntry.references - ) - } - } - - // at the end of recommendation acceptance, report user decisions and clear recommendations. - RecommendationHandler.instance.reportUserDecisions(acceptanceEntry.acceptIndex) -} diff --git a/packages/core/src/codewhisperer/commands/onInlineAcceptance.ts b/packages/core/src/codewhisperer/commands/onInlineAcceptance.ts deleted file mode 100644 index 50af478ba57..00000000000 --- a/packages/core/src/codewhisperer/commands/onInlineAcceptance.ts +++ /dev/null @@ -1,146 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import * as CodeWhispererConstants from '../models/constants' -import { vsCodeState, OnRecommendationAcceptanceEntry } from '../models/model' -import { runtimeLanguageContext } from '../util/runtimeLanguageContext' -import { CodeWhispererTracker } from '../tracker/codewhispererTracker' -import { CodeWhispererCodeCoverageTracker } from '../tracker/codewhispererCodeCoverageTracker' -import { getLogger } from '../../shared/logger/logger' -import { RecommendationHandler } from '../service/recommendationHandler' -import { sleep } from '../../shared/utilities/timeoutUtils' -import { handleExtraBrackets } from '../util/closingBracketUtil' -import { Commands } from '../../shared/vscode/commands2' -import { isInlineCompletionEnabled } from '../util/commonUtil' -import { ExtContext } from '../../shared/extensions' -import { onAcceptance } from './onAcceptance' -import * as codewhispererClient from '../client/codewhisperer' -import { - CodewhispererCompletionType, - CodewhispererLanguage, - CodewhispererTriggerType, -} from '../../shared/telemetry/telemetry.gen' -import { ReferenceLogViewProvider } from '../service/referenceLogViewProvider' -import { ReferenceHoverProvider } from '../service/referenceHoverProvider' -import { ImportAdderProvider } from '../service/importAdderProvider' -import { session } from '../util/codeWhispererSession' -import path from 'path' -import { RecommendationService } from '../service/recommendationService' -import { Container } from '../service/serviceContainer' -import { telemetry } from '../../shared/telemetry/telemetry' -import { TelemetryHelper } from '../util/telemetryHelper' -import { UserWrittenCodeTracker } from '../tracker/userWrittenCodeTracker' - -export const acceptSuggestion = Commands.declare( - 'aws.amazonq.accept', - (context: ExtContext) => - async ( - range: vscode.Range, - effectiveRange: vscode.Range, - acceptIndex: number, - recommendation: string, - requestId: string, - sessionId: string, - triggerType: CodewhispererTriggerType, - completionType: CodewhispererCompletionType, - language: CodewhispererLanguage, - references: codewhispererClient.References - ) => { - telemetry.record({ - traceId: TelemetryHelper.instance.traceId, - }) - - RecommendationService.instance.incrementAcceptedCount() - const editor = vscode.window.activeTextEditor - await Container.instance.lineAnnotationController.refresh(editor, 'codewhisperer') - const onAcceptanceFunc = isInlineCompletionEnabled() ? onInlineAcceptance : onAcceptance - await onAcceptanceFunc({ - editor, - range, - effectiveRange, - acceptIndex, - recommendation, - requestId, - sessionId, - triggerType, - completionType, - language, - references, - }) - } -) -/** - * This function is called when user accepts a intelliSense suggestion or an inline suggestion - */ -export async function onInlineAcceptance(acceptanceEntry: OnRecommendationAcceptanceEntry) { - RecommendationHandler.instance.cancelPaginatedRequest() - RecommendationHandler.instance.disposeInlineCompletion() - - if (acceptanceEntry.editor) { - await sleep(CodeWhispererConstants.vsCodeCursorUpdateDelay) - const languageContext = runtimeLanguageContext.getLanguageContext( - acceptanceEntry.editor.document.languageId, - path.extname(acceptanceEntry.editor.document.fileName) - ) - const start = acceptanceEntry.range.start - const end = acceptanceEntry.editor.selection.active - - vsCodeState.isCodeWhispererEditing = true - /** - * Mitigation to right context handling mainly for auto closing bracket use case - */ - try { - // Do not handle extra bracket if there is a right context merge - if (acceptanceEntry.recommendation === session.recommendations[acceptanceEntry.acceptIndex].content) { - await handleExtraBrackets(acceptanceEntry.editor, end, acceptanceEntry.effectiveRange.start) - } - await ImportAdderProvider.instance.onAcceptRecommendation( - acceptanceEntry.editor, - session.recommendations[acceptanceEntry.acceptIndex], - start.line - ) - } catch (error) { - getLogger().error(`${error} in handling extra brackets or imports`) - } finally { - vsCodeState.isCodeWhispererEditing = false - } - - CodeWhispererTracker.getTracker().enqueue({ - time: new Date(), - fileUrl: acceptanceEntry.editor.document.uri, - originalString: acceptanceEntry.editor.document.getText(new vscode.Range(start, end)), - startPosition: start, - endPosition: end, - requestId: acceptanceEntry.requestId, - sessionId: acceptanceEntry.sessionId, - index: acceptanceEntry.acceptIndex, - triggerType: acceptanceEntry.triggerType, - completionType: acceptanceEntry.completionType, - language: languageContext.language, - }) - const insertedCoderange = new vscode.Range(start, end) - CodeWhispererCodeCoverageTracker.getTracker(languageContext.language)?.countAcceptedTokens( - insertedCoderange, - acceptanceEntry.editor.document.getText(insertedCoderange), - acceptanceEntry.editor.document.fileName - ) - UserWrittenCodeTracker.instance.onQFinishesEdits() - if (acceptanceEntry.references !== undefined) { - const referenceLog = ReferenceLogViewProvider.getReferenceLog( - acceptanceEntry.recommendation, - acceptanceEntry.references, - acceptanceEntry.editor - ) - ReferenceLogViewProvider.instance.addReferenceLog(referenceLog) - ReferenceHoverProvider.instance.addCodeReferences( - acceptanceEntry.recommendation, - acceptanceEntry.references - ) - } - - RecommendationHandler.instance.reportUserDecisions(acceptanceEntry.acceptIndex) - } -} diff --git a/packages/core/src/codewhisperer/commands/startSecurityScan.ts b/packages/core/src/codewhisperer/commands/startSecurityScan.ts index d04fe6effc3..344f3b0eb58 100644 --- a/packages/core/src/codewhisperer/commands/startSecurityScan.ts +++ b/packages/core/src/codewhisperer/commands/startSecurityScan.ts @@ -135,7 +135,7 @@ export async function startSecurityScan( result: 'Succeeded', codewhispererCodeScanTotalIssues: 0, codewhispererCodeScanIssuesWithFixes: 0, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, codewhispererCodeScanScope: scope, source: initiatedByChat ? 'chat' : 'menu', } diff --git a/packages/core/src/codewhisperer/index.ts b/packages/core/src/codewhisperer/index.ts index 4235ae28668..7f1cd3b0880 100644 --- a/packages/core/src/codewhisperer/index.ts +++ b/packages/core/src/codewhisperer/index.ts @@ -9,13 +9,6 @@ export * from './models/model' export * from './models/constants' export * from './commands/basicCommands' export * from './commands/types' -export { - AutotriggerState, - EndState, - ManualtriggerState, - PressTabState, - TryMoreExState, -} from './views/lineAnnotationController' export type { TransformationProgressUpdate, TransformationStep, @@ -43,7 +36,7 @@ export { codeWhispererClient, } from './client/codewhisperer' export { listCodeWhispererCommands, listCodeWhispererCommandsId } from './ui/statusBarMenu' -export { refreshStatusBar, CodeWhispererStatusBar, InlineCompletionService } from './service/inlineCompletionService' +export { refreshStatusBar, CodeWhispererStatusBarManager } from './service/statusBar' export { SecurityIssueHoverProvider } from './service/securityIssueHoverProvider' export { SecurityIssueCodeActionProvider } from './service/securityIssueCodeActionProvider' export { @@ -53,48 +46,30 @@ export { IssueItem, SeverityItem, } from './service/securityIssueTreeViewProvider' -export { invokeRecommendation } from './commands/invokeRecommendation' -export { onAcceptance } from './commands/onAcceptance' export { CodeWhispererTracker } from './tracker/codewhispererTracker' -export { RecommendationHandler } from './service/recommendationHandler' export { CodeWhispererUserGroupSettings } from './util/userGroupUtil' export { session } from './util/codeWhispererSession' -export { onInlineAcceptance } from './commands/onInlineAcceptance' export { stopTransformByQ } from './commands/startTransformByQ' -export { getCompletionItems, getCompletionItem, getLabel } from './service/completionProvider' export { featureDefinitions, FeatureConfigProvider } from '../shared/featureConfig' export { ReferenceInlineProvider } from './service/referenceInlineProvider' export { ReferenceHoverProvider } from './service/referenceHoverProvider' -export { CWInlineCompletionItemProvider } from './service/inlineCompletionItemProvider' -export { RecommendationService } from './service/recommendationService' -export { ClassifierTrigger } from './service/classifierTrigger' -export { DocumentChangedSource, KeyStrokeHandler, DefaultDocumentChangedType } from './service/keyStrokeHandler' export { ReferenceLogViewProvider } from './service/referenceLogViewProvider' export { ImportAdderProvider } from './service/importAdderProvider' export { LicenseUtil } from './util/licenseUtil' export { SecurityIssueProvider } from './service/securityIssueProvider' export { listScanResults, mapToAggregatedList, pollScanJobStatus } from './service/securityScanHandler' -export { CodeWhispererCodeCoverageTracker } from './tracker/codewhispererCodeCoverageTracker' export { TelemetryHelper } from './util/telemetryHelper' export { LineSelection, LineTracker } from './tracker/lineTracker' -export { BM25Okapi } from './util/supplementalContext/rankBm25' -export { handleExtraBrackets } from './util/closingBracketUtil' export { runtimeLanguageContext, RuntimeLanguageContext } from './util/runtimeLanguageContext' export * as startSecurityScan from './commands/startSecurityScan' -export * from './util/supplementalContext/utgUtils' -export * from './util/supplementalContext/crossFileContextUtil' -export * from './util/editorContext' export * from './util/showSsoPrompt' export * from './util/securityScanLanguageContext' export * from './util/importAdderUtil' -export * from './util/globalStateUtil' export * from './util/zipUtil' export * from './util/diagnosticsUtil' export * from './util/commonUtil' -export * from './util/supplementalContext/codeParsingUtil' -export * from './util/supplementalContext/supplementalContextUtil' export * from './util/codewhispererSettings' -export * as supplementalContextUtil from './util/supplementalContext/supplementalContextUtil' +export * as getStartUrl from './util/getStartUrl' export * from './service/diagnosticsProvider' export * as diagnosticsProvider from './service/diagnosticsProvider' export * from './ui/codeWhispererNodes' @@ -106,6 +81,7 @@ export { baseCustomization, onProfileChangedListener, CustomizationProvider, + notifyNewCustomizations, } from './util/customizationUtil' export { Container } from './service/serviceContainer' export * from './util/gitUtil' diff --git a/packages/core/src/codewhisperer/models/constants.ts b/packages/core/src/codewhisperer/models/constants.ts index 289a89828c3..7839b1c3f8e 100644 --- a/packages/core/src/codewhisperer/models/constants.ts +++ b/packages/core/src/codewhisperer/models/constants.ts @@ -138,7 +138,7 @@ export const runningSecurityScan = 'Reviewing project for code issues...' export const runningFileScan = 'Reviewing current file for code issues...' -export const noSuggestions = 'No suggestions from Amazon Q' +export const noInlineSuggestionsMsg = 'No suggestions from Amazon Q' export const licenseFilter = 'Amazon Q suggestions were filtered due to reference settings' @@ -180,15 +180,9 @@ export const securityScanLearnMoreUri = 'https://docs.aws.amazon.com/amazonq/lat export const identityPoolID = 'us-east-1:70717e99-906f-4add-908c-bd9074a2f5b9' /** - * the interval of the background thread invocation, which is triggered by the timer + * Delay for making requests once the user stops typing. Without a delay, inline suggestions request is triggered every keystroke. */ -export const defaultCheckPeriodMillis = 1000 * 60 * 5 - -// suggestion show delay, in milliseconds -export const suggestionShowDelay = 250 - -// add 200ms more delay on top of inline default 30-50ms -export const inlineSuggestionShowDelay = 200 +export const inlineCompletionsDebounceDelay = 200 export const referenceLog = 'Code Reference Log' diff --git a/packages/core/src/codewhisperer/models/model.ts b/packages/core/src/codewhisperer/models/model.ts index 128d34757fc..b916a8e26dd 100644 --- a/packages/core/src/codewhisperer/models/model.ts +++ b/packages/core/src/codewhisperer/models/model.ts @@ -33,6 +33,10 @@ interface VsCodeState { * Flag indicates whether codewhisperer is doing vscode.TextEditor.edit */ isCodeWhispererEditing: boolean + /** + * Keeps track of whether or not recommendations are currently running + */ + isRecommendationsActive: boolean /** * Timestamp of previous user edit */ @@ -44,6 +48,9 @@ interface VsCodeState { export const vsCodeState: VsCodeState = { isIntelliSenseActive: false, isCodeWhispererEditing: false, + // hack to globally keep track of whether or not recommendations are currently running. This allows us to know + // when recommendations have ran during e2e tests + isRecommendationsActive: false, lastUserModificationTime: 0, isFreeTierLimitReached: false, } diff --git a/packages/core/src/codewhisperer/region/regionProfileManager.ts b/packages/core/src/codewhisperer/region/regionProfileManager.ts index a85a2133d89..2645c573249 100644 --- a/packages/core/src/codewhisperer/region/regionProfileManager.ts +++ b/packages/core/src/codewhisperer/region/regionProfileManager.ts @@ -8,13 +8,6 @@ import { getIcon } from '../../shared/icons' import { DataQuickPickItem } from '../../shared/ui/pickerPrompter' import { CodeWhispererConfig, RegionProfile } from '../models/model' import { showConfirmationMessage } from '../../shared/utilities/messages' -import { - Connection, - isBuilderIdConnection, - isIdcSsoConnection, - isSsoConnection, - SsoConnection, -} from '../../auth/connection' import globals from '../../shared/extensionGlobals' import { once } from '../../shared/utilities/functionUtils' import CodeWhispererUserClient from '../client/codewhispereruserclient' @@ -28,8 +21,10 @@ import { parse } from '@aws-sdk/util-arn-parser' import { isAwsError, ToolkitError } from '../../shared/errors' import { telemetry } from '../../shared/telemetry/telemetry' import { localize } from '../../shared/utilities/vsCodeUtils' +import { IAuthProvider } from '../util/authUtil' import { Commands } from '../../shared/vscode/commands2' import { CachedResource } from '../../shared/utilities/resourceCache' +import { GlobalStatePoller } from '../../shared/globalState' // TODO: is there a better way to manage all endpoint strings in one place? export const defaultServiceConfig: CodeWhispererConfig = { @@ -43,6 +38,9 @@ const endpoints = createConstantMap({ 'eu-central-1': 'https://q.eu-central-1.amazonaws.com/', }) +const getRegionProfiles = () => + globals.globalState.tryGet<{ [label: string]: RegionProfile }>('aws.amazonq.regionProfiles', Object, {}) + /** * 'user' -> users change the profile through Q menu * 'auth' -> users change the profile through webview profile selector page @@ -61,7 +59,6 @@ export class RegionProfileManager { private _activeRegionProfile: RegionProfile | undefined private _onDidChangeRegionProfile = new vscode.EventEmitter() public readonly onDidChangeRegionProfile = this._onDidChangeRegionProfile.event - // Store the last API results (for UI propuse) so we don't need to call service again if doesn't require "latest" result private _profiles: RegionProfile[] = [] @@ -86,22 +83,31 @@ export class RegionProfileManager { } })(this.listRegionProfile.bind(this)) + // This is a poller that handles synchornization of selected region profiles between different IDE windows. + // It checks for changes in global state of region profile, invoking the change handler to switch profiles + public globalStatePoller = GlobalStatePoller.create({ + getState: getRegionProfiles, + changeHandler: async () => { + const profile = this.loadPersistedRegionProfiles() + void this._switchRegionProfile(profile[this.authProvider.profileName], 'reload') + }, + pollIntervalInMs: 2000, + }) + get activeRegionProfile() { - const conn = this.connectionProvider() - if (isBuilderIdConnection(conn)) { + if (this.authProvider.isBuilderIdConnection()) { return undefined } return this._activeRegionProfile } get clientConfig(): CodeWhispererConfig { - const conn = this.connectionProvider() - if (!conn) { + if (!this.authProvider.isConnected()) { throw new ToolkitError('trying to get client configuration without credential') } // builder id should simply use default IAD - if (isBuilderIdConnection(conn)) { + if (this.authProvider.isBuilderIdConnection()) { return defaultServiceConfig } @@ -129,7 +135,7 @@ export class RegionProfileManager { return this._profiles } - constructor(private readonly connectionProvider: () => Connection | undefined) {} + constructor(private readonly authProvider: IAuthProvider) {} async getProfiles(): Promise { return this.cache.getResource() @@ -138,15 +144,14 @@ export class RegionProfileManager { async listRegionProfile(): Promise { this._profiles = [] - const conn = this.connectionProvider() - if (conn === undefined || !isSsoConnection(conn)) { + if (!this.authProvider.isConnected() || !this.authProvider.isSsoSession()) { return [] } const availableProfiles: RegionProfile[] = [] const failedRegions: string[] = [] for (const [region, endpoint] of endpoints.entries()) { - const client = await this._createQClient(region, endpoint, conn as SsoConnection) + const client = await this._createQClient(region, endpoint) const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) => client.listAvailableProfiles(request).promise() const request: CodeWhispererUserClient.ListAvailableProfilesRequest = {} @@ -195,8 +200,7 @@ export class RegionProfileManager { } async switchRegionProfile(regionProfile: RegionProfile | undefined, source: ProfileSwitchIntent) { - const conn = this.connectionProvider() - if (conn === undefined || !isIdcSsoConnection(conn)) { + if (!this.authProvider.isConnected() || !this.authProvider.isIdcConnection()) { return } @@ -204,9 +208,6 @@ export class RegionProfileManager { return } - // TODO: make it typesafe - const ssoConn = this.connectionProvider() as SsoConnection - // only prompt to users when users switch from A profile to B profile if (source !== 'customization' && this.activeRegionProfile !== undefined && regionProfile !== undefined) { const response = await showConfirmationMessage({ @@ -224,9 +225,9 @@ export class RegionProfileManager { telemetry.amazonq_didSelectProfile.emit({ source: source, amazonQProfileRegion: this.activeRegionProfile?.region ?? 'not-set', - ssoRegion: ssoConn.ssoRegion, + ssoRegion: this.authProvider.connection?.region, result: 'Cancelled', - credentialStartUrl: ssoConn.startUrl, + credentialStartUrl: this.authProvider.connection?.startUrl, profileCount: this.profiles.length, }) return @@ -243,9 +244,9 @@ export class RegionProfileManager { telemetry.amazonq_didSelectProfile.emit({ source: source, amazonQProfileRegion: regionProfile?.region ?? 'not-set', - ssoRegion: ssoConn.ssoRegion, + ssoRegion: this.authProvider.connection?.region, result: 'Succeeded', - credentialStartUrl: ssoConn.startUrl, + credentialStartUrl: this.authProvider.connection?.startUrl, profileCount: this.profiles.length, }) } @@ -254,6 +255,10 @@ export class RegionProfileManager { } private async _switchRegionProfile(regionProfile: RegionProfile | undefined, source: ProfileSwitchIntent) { + if (this._activeRegionProfile?.arn === regionProfile?.arn) { + return + } + this._activeRegionProfile = regionProfile this._onDidChangeRegionProfile.fire({ @@ -273,18 +278,30 @@ export class RegionProfileManager { } restoreProfileSelection = once(async () => { - const conn = this.connectionProvider() - if (conn) { - await this.restoreRegionProfile(conn) + if (this.authProvider.isConnected()) { + await this.restoreRegionProfile() } }) - // Note: should be called after [AuthUtil.instance.conn] returns non null - async restoreRegionProfile(conn: Connection) { - const previousSelected = this.loadPersistedRegionProfle()[conn.id] || undefined - if (!previousSelected) { + // Note: should be called after [this.authProvider.isConnected()] returns non null + async restoreRegionProfile() { + const profiles = this.loadPersistedRegionProfiles() + if (!profiles || Object.keys(profiles).length === 0) { + return + } + + let previousSelected = profiles[this.authProvider.profileName] + + // If no profile matches auth profileName and there are multiple profiles, return so user can select + if (!previousSelected && Object.keys(profiles).length > 1) { return } + + // If no profile matches auth profileName but there's only one profile, use that one + if (!previousSelected && Object.keys(profiles).length === 1) { + previousSelected = Object.values(profiles)[0] + } + // cross-validation this.getProfiles() .then(async (profiles) => { @@ -315,32 +332,20 @@ export class RegionProfileManager { await this.switchRegionProfile(previousSelected, 'reload') } - private loadPersistedRegionProfle(): { [label: string]: RegionProfile } { - const previousPersistedState = globals.globalState.tryGet<{ [label: string]: RegionProfile }>( - 'aws.amazonq.regionProfiles', - Object, - {} - ) - - return previousPersistedState + public loadPersistedRegionProfiles(): { [label: string]: RegionProfile } { + return getRegionProfiles() } async persistSelectRegionProfile() { - const conn = this.connectionProvider() - // default has empty arn and shouldn't be persisted because it's just a fallback - if (!conn || this.activeRegionProfile === undefined) { + if (!this.authProvider.isConnected() || this.activeRegionProfile === undefined) { return } // persist connectionId to profileArn - const previousPersistedState = globals.globalState.tryGet<{ [label: string]: RegionProfile }>( - 'aws.amazonq.regionProfiles', - Object, - {} - ) + const previousPersistedState = getRegionProfiles() - previousPersistedState[conn.id] = this.activeRegionProfile + previousPersistedState[this.authProvider.profileName] = this.activeRegionProfile await globals.globalState.update('aws.amazonq.regionProfiles', previousPersistedState) } @@ -387,7 +392,7 @@ export class RegionProfileManager { this._activeRegionProfile = undefined } - const profiles = this.loadPersistedRegionProfle() + const profiles = this.loadPersistedRegionProfiles() const updatedProfiles = Object.fromEntries( Object.entries(profiles).filter(([connId, profile]) => profile.arn !== arn) ) @@ -395,27 +400,32 @@ export class RegionProfileManager { } } - // Should be called on connection changed in case users change to a differnet connection and use the wrong resultset. + requireProfileSelection(): boolean { + if (this.authProvider.isBuilderIdConnection()) { + return false + } + return this.authProvider.isIdcConnection() && this.activeRegionProfile === undefined + } + async clearCache() { await this.cache.clearCache() } // TODO: Should maintain sdk client in a better way async createQClient(profile: RegionProfile): Promise { - const conn = this.connectionProvider() - if (conn === undefined || !isSsoConnection(conn)) { + if (!this.authProvider.isConnected() || !this.authProvider.isSsoSession()) { throw new Error('No valid SSO connection') } const endpoint = endpoints.get(profile.region) if (!endpoint) { throw new Error(`trying to initiatize Q client with unrecognizable region ${profile.region}`) } - return this._createQClient(profile.region, endpoint, conn) + return this._createQClient(profile.region, endpoint) } // Visible for testing only, do not use this directly, please use createQClient(profile) - async _createQClient(region: string, endpoint: string, conn: SsoConnection): Promise { - const token = (await conn.getToken()).accessToken + async _createQClient(region: string, endpoint: string): Promise { + const token = await this.authProvider.getToken() const serviceOption: ServiceOptions = { apiConfig: userApiConfig, region: region, diff --git a/packages/core/src/codewhisperer/region/utils.ts b/packages/core/src/codewhisperer/region/utils.ts index dd988f74a30..fb768e3b710 100644 --- a/packages/core/src/codewhisperer/region/utils.ts +++ b/packages/core/src/codewhisperer/region/utils.ts @@ -7,7 +7,6 @@ const localize = nls.loadMessageBundle() import { AmazonQPromptSettings } from '../../shared/settings' import { telemetry } from '../../shared/telemetry/telemetry' import vscode from 'vscode' -import { selectRegionProfileCommand } from '../commands/basicCommands' import { placeholder } from '../../shared/vscode/commands2' import { toastMessage } from '../commands/types' @@ -36,7 +35,7 @@ export async function notifySelectDeveloperProfile() { if (resp === selectProfile) { // Show Profile telemetry.record({ action: 'select' }) - void selectRegionProfileCommand.execute(placeholder, toastMessage) + void vscode.commands.executeCommand('aws.amazonq.selectRegionProfile', placeholder, toastMessage) } else if (resp === dontShowAgain) { telemetry.record({ action: 'dontShowAgain' }) await settings.disablePrompt(suppressId) diff --git a/packages/core/src/codewhisperer/service/classifierTrigger.ts b/packages/core/src/codewhisperer/service/classifierTrigger.ts deleted file mode 100644 index 842d5312e68..00000000000 --- a/packages/core/src/codewhisperer/service/classifierTrigger.ts +++ /dev/null @@ -1,609 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import os from 'os' -import * as vscode from 'vscode' -import { CodewhispererAutomatedTriggerType } from '../../shared/telemetry/telemetry' -import { extractContextForCodeWhisperer } from '../util/editorContext' -import { TelemetryHelper } from '../util/telemetryHelper' -import { ProgrammingLanguage } from '../client/codewhispereruserclient' - -interface normalizedCoefficients { - readonly lineNum: number - readonly lenLeftCur: number - readonly lenLeftPrev: number - readonly lenRight: number -} -/* - uses ML classifier to determine if user input should trigger CWSPR service - */ -export class ClassifierTrigger { - static #instance: ClassifierTrigger - - public static get instance() { - return (this.#instance ??= new this()) - } - - // ML classifier trigger threshold - private triggerThreshold = 0.43 - - // ML classifier coefficients - // os coefficient - private osCoefficientMap: Readonly> = { - 'Mac OS X': -0.1552, - 'Windows 10': -0.0238, - Windows: 0.0412, - win32: -0.0559, - } - - // trigger type coefficient - private triggerTypeCoefficientMap: Readonly> = { - SpecialCharacters: 0.0209, - Enter: 0.2853, - } - - private languageCoefficientMap: Readonly> = { - java: -0.4622, - javascript: -0.4688, - python: -0.3052, - typescript: -0.6084, - tsx: -0.6084, - jsx: -0.4688, - shell: -0.4718, - ruby: -0.7356, - sql: -0.4937, - rust: -0.4309, - kotlin: -0.4739, - php: -0.3917, - csharp: -0.3475, - go: -0.3504, - scala: -0.534, - cpp: -0.1734, - json: 0, - yaml: -0.3, - tf: -0.55, - } - - // other metadata coefficient - private lineNumCoefficient = -0.0416 - private lengthOfLeftCurrentCoefficient = -1.1747 - private lengthOfLeftPrevCoefficient = 0.4033 - private lengthOfRightCoefficient = -0.3321 - private prevDecisionAcceptCoefficient = 0.5397 - private prevDecisionRejectCoefficient = -0.1656 - private prevDecisionOtherCoefficient = 0 - private ideVscode = -0.1905 - private lengthLeft0To5 = -0.8756 - private lengthLeft5To10 = -0.5463 - private lengthLeft10To20 = -0.4081 - private lengthLeft20To30 = -0.3272 - private lengthLeft30To40 = -0.2442 - private lengthLeft40To50 = -0.1471 - - // intercept of logistic regression classifier - private intercept = 0.3738713 - - private maxx: normalizedCoefficients = { - lineNum: 4631.0, - lenLeftCur: 157.0, - lenLeftPrev: 176.0, - lenRight: 10239.0, - } - - private minn: normalizedCoefficients = { - lineNum: 0.0, - lenLeftCur: 0.0, - lenLeftPrev: 0.0, - lenRight: 0.0, - } - - // character and keywords coefficient - private charCoefficient: Readonly> = { - throw: 1.5868, - ';': -1.268, - any: -1.1565, - '7': -1.1347, - false: -1.1307, - nil: -1.0653, - elif: 1.0122, - '9': -1.0098, - pass: -1.0058, - True: -1.0002, - False: -0.9434, - '6': -0.9222, - true: -0.9142, - None: -0.9027, - '8': -0.9013, - break: -0.8475, - '}': -0.847, - '5': -0.8414, - '4': -0.8197, - '1': -0.8085, - '\\': -0.8019, - static: -0.7748, - '0': -0.77, - end: -0.7617, - '(': 0.7239, - '/': -0.7104, - where: -0.6981, - readonly: -0.6741, - async: -0.6723, - '3': -0.654, - continue: -0.6413, - struct: -0.64, - try: -0.6369, - float: -0.6341, - using: 0.6079, - '@': 0.6016, - '|': 0.5993, - impl: 0.5808, - private: -0.5746, - for: 0.5741, - '2': -0.5634, - let: -0.5187, - foreach: 0.5186, - select: -0.5148, - export: -0.5, - mut: -0.4921, - ')': -0.463, - ']': -0.4611, - when: 0.4602, - virtual: -0.4583, - extern: -0.4465, - catch: 0.4446, - new: 0.4394, - val: -0.4339, - map: 0.4284, - case: 0.4271, - throws: 0.4221, - null: -0.4197, - protected: -0.4133, - q: 0.4125, - except: 0.4115, - ': ': 0.4072, - '^': -0.407, - ' ': 0.4066, - $: 0.3981, - this: 0.3962, - switch: 0.3947, - '*': -0.3931, - module: 0.3912, - array: 0.385, - '=': 0.3828, - p: 0.3728, - ON: 0.3708, - '`': 0.3693, - u: 0.3658, - a: 0.3654, - require: 0.3646, - '>': -0.3644, - const: -0.3476, - o: 0.3423, - sizeof: 0.3416, - object: 0.3362, - w: 0.3345, - print: 0.3344, - range: 0.3336, - if: 0.3324, - abstract: -0.3293, - var: -0.3239, - i: 0.321, - while: 0.3138, - J: 0.3137, - c: 0.3118, - await: -0.3072, - from: 0.3057, - f: 0.302, - echo: 0.2995, - '#': 0.2984, - e: 0.2962, - r: 0.2925, - mod: 0.2893, - loop: 0.2874, - t: 0.2832, - '~': 0.282, - final: -0.2816, - del: 0.2785, - override: -0.2746, - ref: -0.2737, - h: 0.2693, - m: 0.2681, - '{': 0.2674, - implements: 0.2672, - inline: -0.2642, - match: 0.2613, - with: -0.261, - x: 0.2597, - namespace: -0.2596, - operator: 0.2573, - double: -0.2563, - source: -0.2482, - import: -0.2419, - NULL: -0.2399, - l: 0.239, - or: 0.2378, - s: 0.2366, - then: 0.2354, - W: 0.2354, - y: 0.2333, - local: 0.2288, - is: 0.2282, - n: 0.2254, - '+': -0.2251, - G: 0.223, - public: -0.2229, - WHERE: 0.2224, - list: 0.2204, - Q: 0.2204, - '[': 0.2136, - VALUES: 0.2134, - H: 0.2105, - g: 0.2094, - else: -0.208, - bool: -0.2066, - long: -0.2059, - R: 0.2025, - S: 0.2021, - d: 0.2003, - V: 0.1974, - K: -0.1961, - '<': 0.1958, - debugger: -0.1929, - NOT: -0.1911, - b: 0.1907, - boolean: -0.1891, - z: -0.1866, - LIKE: -0.1793, - raise: 0.1782, - L: 0.1768, - fn: 0.176, - delete: 0.1714, - unsigned: -0.1675, - auto: -0.1648, - finally: 0.1616, - k: 0.1599, - as: 0.156, - instanceof: 0.1558, - '&': 0.1554, - E: 0.1551, - M: 0.1542, - I: 0.1503, - Y: 0.1493, - typeof: 0.1475, - j: 0.1445, - INTO: 0.1442, - IF: 0.1437, - next: 0.1433, - undef: -0.1427, - THEN: -0.1416, - v: 0.1415, - C: 0.1383, - P: 0.1353, - AND: -0.1345, - constructor: 0.1337, - void: -0.1336, - class: -0.1328, - defer: 0.1316, - begin: 0.1306, - FROM: -0.1304, - SET: 0.1291, - decimal: -0.1278, - friend: 0.1277, - SELECT: -0.1265, - event: 0.1259, - lambda: 0.1253, - enum: 0.1215, - A: 0.121, - lock: 0.1187, - ensure: 0.1184, - '%': 0.1177, - isset: 0.1175, - O: 0.1174, - '.': 0.1146, - UNION: -0.1145, - alias: -0.1129, - template: -0.1102, - WHEN: 0.1093, - rescue: 0.1083, - DISTINCT: -0.1074, - trait: -0.1073, - D: 0.1062, - in: 0.1045, - internal: -0.1029, - ',': 0.1027, - static_cast: 0.1016, - do: -0.1005, - OR: 0.1003, - AS: -0.1001, - interface: 0.0996, - super: 0.0989, - B: 0.0963, - U: 0.0962, - T: 0.0943, - CALL: -0.0918, - BETWEEN: -0.0915, - N: 0.0897, - yield: 0.0867, - done: -0.0857, - string: -0.0837, - out: -0.0831, - volatile: -0.0819, - retry: 0.0816, - '?': -0.0796, - number: -0.0791, - short: 0.0787, - sealed: -0.0776, - package: 0.0765, - OPEN: -0.0756, - base: 0.0735, - and: 0.0729, - exit: 0.0726, - _: 0.0721, - keyof: -0.072, - def: 0.0713, - crate: -0.0706, - '-': -0.07, - FUNCTION: 0.0692, - declare: -0.0678, - include: 0.0671, - COUNT: -0.0669, - INDEX: -0.0666, - CLOSE: -0.0651, - fi: -0.0644, - uint: 0.0624, - params: 0.0575, - HAVING: 0.0575, - byte: -0.0575, - clone: -0.0552, - char: -0.054, - func: 0.0538, - never: -0.053, - unset: -0.0524, - unless: -0.051, - esac: -0.0509, - shift: -0.0507, - require_once: 0.0486, - ELSE: -0.0477, - extends: 0.0461, - elseif: 0.0452, - mutable: -0.0451, - asm: 0.0449, - '!': 0.0446, - LIMIT: 0.0444, - ushort: -0.0438, - '"': -0.0433, - Z: 0.0431, - exec: -0.0431, - IS: -0.0429, - DECLARE: -0.0425, - __LINE__: -0.0424, - BEGIN: -0.0418, - typedef: 0.0414, - EXIT: -0.0412, - "'": 0.041, - function: -0.0393, - dyn: -0.039, - wchar_t: -0.0388, - unique: -0.0383, - include_once: 0.0367, - stackalloc: 0.0359, - RETURN: -0.0356, - const_cast: 0.035, - MAX: 0.0341, - assert: -0.0331, - JOIN: -0.0328, - use: 0.0318, - GET: 0.0317, - VIEW: 0.0314, - move: 0.0308, - typename: 0.0308, - die: 0.0305, - asserts: -0.0304, - reinterpret_cast: -0.0302, - USING: -0.0289, - elsif: -0.0285, - FIRST: -0.028, - self: -0.0278, - RETURNING: -0.0278, - symbol: -0.0273, - OFFSET: 0.0263, - bigint: 0.0253, - register: -0.0237, - union: -0.0227, - return: -0.0227, - until: -0.0224, - endfor: -0.0213, - implicit: -0.021, - LOOP: 0.0195, - pub: 0.0182, - global: 0.0179, - EXCEPTION: 0.0175, - delegate: 0.0173, - signed: -0.0163, - FOR: 0.0156, - unsafe: 0.014, - NEXT: -0.0133, - IN: 0.0129, - MIN: -0.0123, - go: -0.0112, - type: -0.0109, - explicit: -0.0107, - eval: -0.0104, - int: -0.0099, - CASE: -0.0096, - END: 0.0084, - UPDATE: 0.0074, - default: 0.0072, - chan: 0.0068, - fixed: 0.0066, - not: -0.0052, - X: -0.0047, - endforeach: 0.0031, - goto: 0.0028, - empty: 0.0022, - checked: 0.0012, - F: -0.001, - } - - public getThreshold() { - return this.triggerThreshold - } - - public recordClassifierResultForManualTrigger(editor: vscode.TextEditor) { - this.shouldTriggerFromClassifier(undefined, editor, undefined, true) - } - - public recordClassifierResultForAutoTrigger( - editor: vscode.TextEditor, - triggerType?: CodewhispererAutomatedTriggerType, - event?: vscode.TextDocumentChangeEvent - ) { - if (!triggerType) { - return - } - this.shouldTriggerFromClassifier(event, editor, triggerType, true) - } - - public shouldTriggerFromClassifier( - event: vscode.TextDocumentChangeEvent | undefined, - editor: vscode.TextEditor, - autoTriggerType: string | undefined, - shouldRecordResult: boolean = false - ): boolean { - const fileContext = extractContextForCodeWhisperer(editor) - const osPlatform = this.normalizeOsName(os.platform(), os.version()) - const char = event ? event.contentChanges[0].text : '' - const lineNum = editor.selection.active.line - const classifierResult = this.getClassifierResult( - fileContext.leftFileContent, - fileContext.rightFileContent, - osPlatform, - autoTriggerType, - char, - lineNum, - fileContext.programmingLanguage - ) - - const threshold = this.getThreshold() - - const shouldTrigger = classifierResult > threshold - if (shouldRecordResult) { - TelemetryHelper.instance.setClassifierResult(classifierResult) - TelemetryHelper.instance.setClassifierThreshold(threshold) - } - return shouldTrigger - } - - private getClassifierResult( - leftContext: string, - rightContext: string, - os: string, - triggerType: string | undefined, - char: string, - lineNum: number, - language: ProgrammingLanguage - ): number { - const leftContextLines = leftContext.split(/\r?\n/) - const leftContextAtCurrentLine = leftContextLines[leftContextLines.length - 1] - const tokens = leftContextAtCurrentLine.trim().split(' ') - let keyword = '' - const lastToken = tokens[tokens.length - 1] - if (lastToken && lastToken.length > 1) { - keyword = lastToken - } - const lengthOfLeftCurrent = leftContextLines[leftContextLines.length - 1].length - const lengthOfLeftPrev = leftContextLines[leftContextLines.length - 2]?.length ?? 0 - const lengthOfRight = rightContext.trim().length - - const triggerTypeCoefficient: number = this.triggerTypeCoefficientMap[triggerType || ''] ?? 0 - const osCoefficient: number = this.osCoefficientMap[os] ?? 0 - const charCoefficient: number = this.charCoefficient[char] ?? 0 - const keyWordCoefficient: number = this.charCoefficient[keyword] ?? 0 - const ideCoefficient = this.ideVscode - - const previousDecision = TelemetryHelper.instance.getLastTriggerDecisionForClassifier() - const languageCoefficients = Object.values(this.languageCoefficientMap) - const avrgCoefficient = - languageCoefficients.length > 0 - ? languageCoefficients.reduce((a, b) => a + b) / languageCoefficients.length - : 0 - const languageCoefficient = this.languageCoefficientMap[language.languageName] ?? avrgCoefficient - - let previousDecisionCoefficient = 0 - if (previousDecision === 'Accept') { - previousDecisionCoefficient = this.prevDecisionAcceptCoefficient - } else if (previousDecision === 'Reject') { - previousDecisionCoefficient = this.prevDecisionRejectCoefficient - } else if (previousDecision === 'Discard' || previousDecision === 'Empty') { - previousDecisionCoefficient = this.prevDecisionOtherCoefficient - } - - let leftContextLengthCoefficient = 0 - if (leftContext.length >= 0 && leftContext.length < 5) { - leftContextLengthCoefficient = this.lengthLeft0To5 - } else if (leftContext.length >= 5 && leftContext.length < 10) { - leftContextLengthCoefficient = this.lengthLeft5To10 - } else if (leftContext.length >= 10 && leftContext.length < 20) { - leftContextLengthCoefficient = this.lengthLeft10To20 - } else if (leftContext.length >= 20 && leftContext.length < 30) { - leftContextLengthCoefficient = this.lengthLeft20To30 - } else if (leftContext.length >= 30 && leftContext.length < 40) { - leftContextLengthCoefficient = this.lengthLeft30To40 - } else if (leftContext.length >= 40 && leftContext.length < 50) { - leftContextLengthCoefficient = this.lengthLeft40To50 - } - - const result = - (this.lengthOfRightCoefficient * (lengthOfRight - this.minn.lenRight)) / - (this.maxx.lenRight - this.minn.lenRight) + - (this.lengthOfLeftCurrentCoefficient * (lengthOfLeftCurrent - this.minn.lenLeftCur)) / - (this.maxx.lenLeftCur - this.minn.lenLeftCur) + - (this.lengthOfLeftPrevCoefficient * (lengthOfLeftPrev - this.minn.lenLeftPrev)) / - (this.maxx.lenLeftPrev - this.minn.lenLeftPrev) + - (this.lineNumCoefficient * (lineNum - this.minn.lineNum)) / (this.maxx.lineNum - this.minn.lineNum) + - osCoefficient + - triggerTypeCoefficient + - charCoefficient + - keyWordCoefficient + - ideCoefficient + - this.intercept + - previousDecisionCoefficient + - languageCoefficient + - leftContextLengthCoefficient - - return sigmoid(result) - } - - private normalizeOsName(name: string, version: string | undefined): string { - const lowercaseName = name.toLowerCase() - if (lowercaseName.includes('windows')) { - if (!version) { - return 'Windows' - } else if (version.includes('Windows NT 10') || version.startsWith('10')) { - return 'Windows 10' - } else if (version.includes('6.1')) { - return 'Windows 7' - } else if (version.includes('6.3')) { - return 'Windows 8.1' - } else { - return 'Windows' - } - } else if ( - lowercaseName.includes('macos') || - lowercaseName.includes('mac os') || - lowercaseName.includes('darwin') - ) { - return 'Mac OS X' - } else if (lowercaseName.includes('linux')) { - return 'Linux' - } else { - return name - } - } -} - -const sigmoid = (x: number) => { - return 1 / (1 + Math.exp(-x)) -} diff --git a/packages/core/src/codewhisperer/service/completionProvider.ts b/packages/core/src/codewhisperer/service/completionProvider.ts deleted file mode 100644 index 226d04dec2b..00000000000 --- a/packages/core/src/codewhisperer/service/completionProvider.ts +++ /dev/null @@ -1,77 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import * as CodeWhispererConstants from '../models/constants' -import { runtimeLanguageContext } from '../util/runtimeLanguageContext' -import { Recommendation } from '../client/codewhisperer' -import { LicenseUtil } from '../util/licenseUtil' -import { RecommendationHandler } from './recommendationHandler' -import { session } from '../util/codeWhispererSession' -import path from 'path' -/** - * completion provider for intelliSense popup - */ -export function getCompletionItems(document: vscode.TextDocument, position: vscode.Position) { - const completionItems: vscode.CompletionItem[] = [] - for (const [index, recommendation] of session.recommendations.entries()) { - completionItems.push(getCompletionItem(document, position, recommendation, index)) - session.setSuggestionState(index, 'Showed') - } - return completionItems -} - -export function getCompletionItem( - document: vscode.TextDocument, - position: vscode.Position, - recommendationDetail: Recommendation, - recommendationIndex: number -) { - const start = session.startPos - const range = new vscode.Range(start, start) - const recommendation = recommendationDetail.content - const completionItem = new vscode.CompletionItem(recommendation) - completionItem.insertText = new vscode.SnippetString(recommendation) - completionItem.documentation = new vscode.MarkdownString().appendCodeblock(recommendation, document.languageId) - completionItem.kind = vscode.CompletionItemKind.Method - completionItem.detail = CodeWhispererConstants.completionDetail - completionItem.keepWhitespace = true - completionItem.label = getLabel(recommendation) - completionItem.preselect = true - completionItem.sortText = String(recommendationIndex + 1).padStart(10, '0') - completionItem.range = new vscode.Range(start, position) - const languageContext = runtimeLanguageContext.getLanguageContext( - document.languageId, - path.extname(document.fileName) - ) - let references: typeof recommendationDetail.references - if (recommendationDetail.references !== undefined && recommendationDetail.references.length > 0) { - references = recommendationDetail.references - const licenses = [ - ...new Set(references.map((r) => `[${r.licenseName}](${LicenseUtil.getLicenseHtml(r.licenseName)})`)), - ].join(', ') - completionItem.documentation.appendMarkdown(CodeWhispererConstants.suggestionDetailReferenceText(licenses)) - } - completionItem.command = { - command: 'aws.amazonq.accept', - title: 'On acceptance', - arguments: [ - range, - recommendationIndex, - recommendation, - RecommendationHandler.instance.requestId, - session.sessionId, - session.triggerType, - session.getCompletionType(recommendationIndex), - languageContext.language, - references, - ], - } - return completionItem -} - -export function getLabel(recommendation: string): string { - return recommendation.slice(0, CodeWhispererConstants.labelLength) + '..' -} diff --git a/packages/core/src/codewhisperer/service/inlineCompletionItemProvider.ts b/packages/core/src/codewhisperer/service/inlineCompletionItemProvider.ts deleted file mode 100644 index a6c424c321d..00000000000 --- a/packages/core/src/codewhisperer/service/inlineCompletionItemProvider.ts +++ /dev/null @@ -1,194 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -import vscode, { Position } from 'vscode' -import { getPrefixSuffixOverlap } from '../util/commonUtil' -import { Recommendation } from '../client/codewhisperer' -import { session } from '../util/codeWhispererSession' -import { TelemetryHelper } from '../util/telemetryHelper' -import { runtimeLanguageContext } from '../util/runtimeLanguageContext' -import { ReferenceInlineProvider } from './referenceInlineProvider' -import { ImportAdderProvider } from './importAdderProvider' -import { application } from '../util/codeWhispererApplication' -import path from 'path' -import { UserWrittenCodeTracker } from '../tracker/userWrittenCodeTracker' - -export class CWInlineCompletionItemProvider implements vscode.InlineCompletionItemProvider { - private activeItemIndex: number | undefined - private nextMove: number - private recommendations: Recommendation[] - private requestId: string - private startPos: Position - private nextToken: string - - private _onDidShow: vscode.EventEmitter = new vscode.EventEmitter() - public readonly onDidShow: vscode.Event = this._onDidShow.event - - public constructor( - itemIndex: number | undefined, - firstMove: number, - recommendations: Recommendation[], - requestId: string, - startPos: Position, - nextToken: string - ) { - this.activeItemIndex = itemIndex - this.nextMove = firstMove - this.recommendations = recommendations - this.requestId = requestId - this.startPos = startPos - this.nextToken = nextToken - } - - get getActiveItemIndex() { - return this.activeItemIndex - } - - public clearActiveItemIndex() { - this.activeItemIndex = undefined - } - - // iterate suggestions and stop at index 0 or index len - 1 - private getIteratingIndexes() { - const len = this.recommendations.length - const startIndex = this.activeItemIndex ? this.activeItemIndex : 0 - const index = [] - if (this.nextMove === 0) { - for (let i = 0; i < len; i++) { - index.push((startIndex + i) % len) - } - } else if (this.nextMove === -1) { - for (let i = startIndex - 1; i >= 0; i--) { - index.push(i) - } - index.push(startIndex) - } else { - for (let i = startIndex + 1; i < len; i++) { - index.push(i) - } - index.push(startIndex) - } - return index - } - - truncateOverlapWithRightContext(document: vscode.TextDocument, suggestion: string, pos: vscode.Position): string { - const trimmedSuggestion = suggestion.trim() - // limit of 5000 for right context matching - const rightContext = document.getText(new vscode.Range(pos, document.positionAt(document.offsetAt(pos) + 5000))) - const overlap = getPrefixSuffixOverlap(trimmedSuggestion, rightContext) - const overlapIndex = suggestion.lastIndexOf(overlap) - if (overlapIndex >= 0) { - const truncated = suggestion.slice(0, overlapIndex) - return truncated.trim().length ? truncated : '' - } else { - return suggestion - } - } - - getInlineCompletionItem( - document: vscode.TextDocument, - r: Recommendation, - start: vscode.Position, - end: vscode.Position, - index: number, - prefix: string - ): vscode.InlineCompletionItem | undefined { - if (!r.content.startsWith(prefix)) { - return undefined - } - const effectiveStart = document.positionAt(document.offsetAt(start) + prefix.length) - const truncatedSuggestion = this.truncateOverlapWithRightContext(document, r.content, end) - if (truncatedSuggestion.length === 0) { - if (session.getSuggestionState(index) !== 'Showed') { - session.setSuggestionState(index, 'Discard') - } - return undefined - } - TelemetryHelper.instance.lastSuggestionInDisplay = truncatedSuggestion - return { - insertText: truncatedSuggestion, - range: new vscode.Range(start, end), - command: { - command: 'aws.amazonq.accept', - title: 'On acceptance', - arguments: [ - new vscode.Range(start, end), - new vscode.Range(effectiveStart, end), - index, - truncatedSuggestion, - this.requestId, - session.sessionId, - session.triggerType, - session.getCompletionType(index), - runtimeLanguageContext.getLanguageContext(document.languageId, path.extname(document.fileName)) - .language, - r.references, - ], - }, - } - } - - // the returned completion items will always only contain one valid item - // this is to trace the current index of visible completion item - // so that reference tracker can show - // This hack can be removed once inlineCompletionAdditions API becomes public - provideInlineCompletionItems( - document: vscode.TextDocument, - position: vscode.Position, - _context: vscode.InlineCompletionContext, - _token: vscode.CancellationToken - ): vscode.ProviderResult { - if (position.line < 0 || position.isBefore(this.startPos)) { - application()._clearCodeWhispererUIListener.fire() - this.activeItemIndex = undefined - return - } - - // There's a chance that the startPos is no longer valid in the current document (e.g. - // when CodeWhisperer got triggered by 'Enter', the original startPos is with indentation - // but then this indentation got removed by VSCode when another new line is inserted, - // before the code reaches here). In such case, we need to update the startPos to be a - // valid one. Otherwise, inline completion which utilizes this position will function - // improperly. - const start = document.validatePosition(this.startPos) - const end = position - const iteratingIndexes = this.getIteratingIndexes() - const prefix = document.getText(new vscode.Range(start, end)).replace(/\r\n/g, '\n') - const matchedCount = session.recommendations.filter( - (r) => r.content.length > 0 && r.content.startsWith(prefix) && r.content !== prefix - ).length - for (const i of iteratingIndexes) { - const r = session.recommendations[i] - const item = this.getInlineCompletionItem(document, r, start, end, i, prefix) - if (item === undefined) { - continue - } - this.activeItemIndex = i - session.setSuggestionState(i, 'Showed') - ReferenceInlineProvider.instance.setInlineReference(this.startPos.line, r.content, r.references) - ImportAdderProvider.instance.onShowRecommendation(document, this.startPos.line, r) - this.nextMove = 0 - TelemetryHelper.instance.setFirstSuggestionShowTime() - session.setPerceivedLatency() - UserWrittenCodeTracker.instance.onQStartsMakingEdits() - this._onDidShow.fire() - if (matchedCount >= 2 || this.nextToken !== '') { - const result = [item] - for (let j = 0; j < matchedCount - 1; j++) { - result.push({ - insertText: `${ - typeof item.insertText === 'string' ? item.insertText : item.insertText.value - }${j}`, - range: item.range, - }) - } - return result - } - return [item] - } - application()._clearCodeWhispererUIListener.fire() - this.activeItemIndex = undefined - return [] - } -} diff --git a/packages/core/src/codewhisperer/service/inlineCompletionService.ts b/packages/core/src/codewhisperer/service/inlineCompletionService.ts deleted file mode 100644 index cc9887adb1f..00000000000 --- a/packages/core/src/codewhisperer/service/inlineCompletionService.ts +++ /dev/null @@ -1,273 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -import * as vscode from 'vscode' -import { CodeSuggestionsState, ConfigurationEntry, GetRecommendationsResponse, vsCodeState } from '../models/model' -import * as CodeWhispererConstants from '../models/constants' -import { DefaultCodeWhispererClient } from '../client/codewhisperer' -import { RecommendationHandler } from './recommendationHandler' -import { CodewhispererAutomatedTriggerType, CodewhispererTriggerType } from '../../shared/telemetry/telemetry' -import { showTimedMessage } from '../../shared/utilities/messages' -import { getLogger } from '../../shared/logger/logger' -import { TelemetryHelper } from '../util/telemetryHelper' -import { AuthUtil } from '../util/authUtil' -import { shared } from '../../shared/utilities/functionUtils' -import { ClassifierTrigger } from './classifierTrigger' -import { getSelectedCustomization } from '../util/customizationUtil' -import { codicon, getIcon } from '../../shared/icons' -import { session } from '../util/codeWhispererSession' -import { noSuggestions } from '../models/constants' -import { Commands } from '../../shared/vscode/commands2' -import { listCodeWhispererCommandsId } from '../ui/statusBarMenu' - -export class InlineCompletionService { - private maxPage = 100 - private statusBar: CodeWhispererStatusBar - private _showRecommendationTimer?: NodeJS.Timer - - constructor(statusBar: CodeWhispererStatusBar = CodeWhispererStatusBar.instance) { - this.statusBar = statusBar - - RecommendationHandler.instance.onDidReceiveRecommendation((e) => { - this.startShowRecommendationTimer() - }) - - CodeSuggestionsState.instance.onDidChangeState(() => { - return this.refreshStatusBar() - }) - } - - static #instance: InlineCompletionService - - public static get instance() { - return (this.#instance ??= new this()) - } - - filePath(): string | undefined { - return RecommendationHandler.instance.documentUri?.fsPath - } - - private sharedTryShowRecommendation = shared( - RecommendationHandler.instance.tryShowRecommendation.bind(RecommendationHandler.instance) - ) - - private startShowRecommendationTimer() { - if (this._showRecommendationTimer) { - clearInterval(this._showRecommendationTimer) - this._showRecommendationTimer = undefined - } - this._showRecommendationTimer = setInterval(() => { - const delay = performance.now() - vsCodeState.lastUserModificationTime - if (delay < CodeWhispererConstants.inlineSuggestionShowDelay) { - return - } - this.sharedTryShowRecommendation() - .catch((e) => { - getLogger().error('tryShowRecommendation failed: %s', (e as Error).message) - }) - .finally(() => { - if (this._showRecommendationTimer) { - clearInterval(this._showRecommendationTimer) - this._showRecommendationTimer = undefined - } - }) - }, CodeWhispererConstants.showRecommendationTimerPollPeriod) - } - - async getPaginatedRecommendation( - client: DefaultCodeWhispererClient, - editor: vscode.TextEditor, - triggerType: CodewhispererTriggerType, - config: ConfigurationEntry, - autoTriggerType?: CodewhispererAutomatedTriggerType, - event?: vscode.TextDocumentChangeEvent - ): Promise { - if (vsCodeState.isCodeWhispererEditing || RecommendationHandler.instance.isSuggestionVisible()) { - return { - result: 'Failed', - errorMessage: 'Amazon Q is already running', - recommendationCount: 0, - } - } - - // Call report user decisions once to report recommendations leftover from last invocation. - RecommendationHandler.instance.reportUserDecisions(-1) - TelemetryHelper.instance.setInvokeSuggestionStartTime() - ClassifierTrigger.instance.recordClassifierResultForAutoTrigger(editor, autoTriggerType, event) - - const triggerChar = event?.contentChanges[0]?.text - if (autoTriggerType === 'SpecialCharacters' && triggerChar) { - TelemetryHelper.instance.setTriggerCharForUserTriggerDecision(triggerChar) - } - const isAutoTrigger = triggerType === 'AutoTrigger' - if (AuthUtil.instance.isConnectionExpired()) { - await AuthUtil.instance.notifyReauthenticate(isAutoTrigger) - return { - result: 'Failed', - errorMessage: 'auth', - recommendationCount: 0, - } - } - - await this.setState('loading') - - RecommendationHandler.instance.checkAndResetCancellationTokens() - RecommendationHandler.instance.documentUri = editor.document.uri - let response: GetRecommendationsResponse = { - result: 'Failed', - errorMessage: undefined, - recommendationCount: 0, - } - try { - let page = 0 - while (page < this.maxPage) { - response = await RecommendationHandler.instance.getRecommendations( - client, - editor, - triggerType, - config, - autoTriggerType, - true, - page - ) - if (RecommendationHandler.instance.checkAndResetCancellationTokens()) { - RecommendationHandler.instance.reportUserDecisions(-1) - await vscode.commands.executeCommand('aws.amazonq.refreshStatusBar') - if (triggerType === 'OnDemand' && session.recommendations.length === 0) { - void showTimedMessage(response.errorMessage ? response.errorMessage : noSuggestions, 2000) - } - return { - result: 'Failed', - errorMessage: 'cancelled', - recommendationCount: 0, - } - } - if (!RecommendationHandler.instance.hasNextToken()) { - break - } - page++ - } - } catch (error) { - getLogger().error(`Error ${error} in getPaginatedRecommendation`) - } - await vscode.commands.executeCommand('aws.amazonq.refreshStatusBar') - if (triggerType === 'OnDemand' && session.recommendations.length === 0) { - void showTimedMessage(response.errorMessage ? response.errorMessage : noSuggestions, 2000) - } - TelemetryHelper.instance.tryRecordClientComponentLatency() - - return { - result: 'Succeeded', - errorMessage: undefined, - recommendationCount: session.recommendations.length, - } - } - - /** Updates the status bar to represent the latest CW state */ - refreshStatusBar() { - if (AuthUtil.instance.isConnectionValid()) { - if (AuthUtil.instance.requireProfileSelection()) { - return this.setState('needsProfile') - } - return this.setState('ok') - } else if (AuthUtil.instance.isConnectionExpired()) { - return this.setState('expired') - } else { - return this.setState('notConnected') - } - } - - private async setState(state: keyof typeof states) { - switch (state) { - case 'loading': { - await this.statusBar.setState('loading') - break - } - case 'ok': { - await this.statusBar.setState('ok', CodeSuggestionsState.instance.isSuggestionsEnabled()) - break - } - case 'expired': { - await this.statusBar.setState('expired') - break - } - case 'notConnected': { - await this.statusBar.setState('notConnected') - break - } - case 'needsProfile': { - await this.statusBar.setState('needsProfile') - break - } - } - } -} - -/** The states that the completion service can be in */ -const states = { - loading: 'loading', - ok: 'ok', - expired: 'expired', - notConnected: 'notConnected', - needsProfile: 'needsProfile', -} as const - -export class CodeWhispererStatusBar { - protected statusBar: vscode.StatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 1) - - static #instance: CodeWhispererStatusBar - static get instance() { - return (this.#instance ??= new this()) - } - - protected constructor() {} - - async setState(state: keyof Omit): Promise - async setState(status: keyof Pick, isSuggestionsEnabled: boolean): Promise - async setState(status: keyof typeof states, isSuggestionsEnabled?: boolean): Promise { - const statusBar = this.statusBar - statusBar.command = listCodeWhispererCommandsId - statusBar.backgroundColor = undefined - - const title = 'Amazon Q' - switch (status) { - case 'loading': { - const selectedCustomization = getSelectedCustomization() - statusBar.text = codicon` ${getIcon('vscode-loading~spin')} ${title}${ - selectedCustomization.arn === '' ? '' : ` | ${selectedCustomization.name}` - }` - break - } - case 'ok': { - const selectedCustomization = getSelectedCustomization() - const icon = isSuggestionsEnabled ? getIcon('vscode-debug-start') : getIcon('vscode-debug-pause') - statusBar.text = codicon`${icon} ${title}${ - selectedCustomization.arn === '' ? '' : ` | ${selectedCustomization.name}` - }` - break - } - - case 'expired': { - statusBar.text = codicon` ${getIcon('vscode-debug-disconnect')} ${title}` - statusBar.backgroundColor = new vscode.ThemeColor('statusBarItem.warningBackground') - break - } - case 'needsProfile': - case 'notConnected': - statusBar.text = codicon` ${getIcon('vscode-chrome-close')} ${title}` - statusBar.backgroundColor = new vscode.ThemeColor('statusBarItem.errorBackground') - break - } - - statusBar.show() - } -} - -/** In this module due to circulare dependency issues */ -export const refreshStatusBar = Commands.declare( - { id: 'aws.amazonq.refreshStatusBar', logging: false }, - () => async () => { - await InlineCompletionService.instance.refreshStatusBar() - } -) diff --git a/packages/core/src/codewhisperer/service/keyStrokeHandler.ts b/packages/core/src/codewhisperer/service/keyStrokeHandler.ts deleted file mode 100644 index 49ef633a98f..00000000000 --- a/packages/core/src/codewhisperer/service/keyStrokeHandler.ts +++ /dev/null @@ -1,267 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import { DefaultCodeWhispererClient } from '../client/codewhisperer' -import * as CodeWhispererConstants from '../models/constants' -import { ConfigurationEntry } from '../models/model' -import { getLogger } from '../../shared/logger/logger' -import { RecommendationHandler } from './recommendationHandler' -import { CodewhispererAutomatedTriggerType } from '../../shared/telemetry/telemetry' -import { getTabSizeSetting } from '../../shared/utilities/editorUtilities' -import { isInlineCompletionEnabled } from '../util/commonUtil' -import { ClassifierTrigger } from './classifierTrigger' -import { extractContextForCodeWhisperer } from '../util/editorContext' -import { RecommendationService } from './recommendationService' - -/** - * This class is for CodeWhisperer auto trigger - */ -export class KeyStrokeHandler { - /** - * Special character which automated triggers codewhisperer - */ - public specialChar: string - /** - * Key stroke count for automated trigger - */ - - private idleTriggerTimer?: NodeJS.Timer - - public lastInvocationTime?: number - - constructor() { - this.specialChar = '' - } - - static #instance: KeyStrokeHandler - - public static get instance() { - return (this.#instance ??= new this()) - } - - public startIdleTimeTriggerTimer( - event: vscode.TextDocumentChangeEvent, - editor: vscode.TextEditor, - client: DefaultCodeWhispererClient, - config: ConfigurationEntry - ) { - if (this.idleTriggerTimer) { - clearInterval(this.idleTriggerTimer) - this.idleTriggerTimer = undefined - } - if (!this.shouldTriggerIdleTime()) { - return - } - this.idleTriggerTimer = setInterval(() => { - const duration = (performance.now() - RecommendationHandler.instance.lastInvocationTime) / 1000 - if (duration < CodeWhispererConstants.invocationTimeIntervalThreshold) { - return - } - - this.invokeAutomatedTrigger('IdleTime', editor, client, config, event) - .catch((e) => { - getLogger().error('invokeAutomatedTrigger failed: %s', (e as Error).message) - }) - .finally(() => { - if (this.idleTriggerTimer) { - clearInterval(this.idleTriggerTimer) - this.idleTriggerTimer = undefined - } - }) - }, CodeWhispererConstants.idleTimerPollPeriod) - } - - public shouldTriggerIdleTime(): boolean { - if (isInlineCompletionEnabled() && RecommendationService.instance.isRunning) { - return false - } - return true - } - - async processKeyStroke( - event: vscode.TextDocumentChangeEvent, - editor: vscode.TextEditor, - client: DefaultCodeWhispererClient, - config: ConfigurationEntry - ): Promise { - try { - if (!config.isAutomatedTriggerEnabled) { - return - } - - // Skip when output channel gains focus and invoke - if (editor.document.languageId === 'Log') { - return - } - - const { rightFileContent } = extractContextForCodeWhisperer(editor) - const rightContextLines = rightFileContent.split(/\r?\n/) - const rightContextAtCurrentLine = rightContextLines[0] - // we do not want to trigger when there is immediate right context on the same line - // with "}" being an exception because of IDE auto-complete - if ( - rightContextAtCurrentLine.length && - !rightContextAtCurrentLine.startsWith(' ') && - rightContextAtCurrentLine.trim() !== '}' && - rightContextAtCurrentLine.trim() !== ')' - ) { - return - } - - let triggerType: CodewhispererAutomatedTriggerType | undefined - const changedSource = new DefaultDocumentChangedType(event.contentChanges).checkChangeSource() - - switch (changedSource) { - case DocumentChangedSource.EnterKey: { - triggerType = 'Enter' - break - } - case DocumentChangedSource.SpecialCharsKey: { - triggerType = 'SpecialCharacters' - break - } - case DocumentChangedSource.RegularKey: { - triggerType = ClassifierTrigger.instance.shouldTriggerFromClassifier(event, editor, triggerType) - ? 'Classifier' - : undefined - break - } - default: { - break - } - } - - if (triggerType) { - await this.invokeAutomatedTrigger(triggerType, editor, client, config, event) - } - } catch (error) { - getLogger().verbose(`Automated Trigger Exception : ${error}`) - } - } - - async invokeAutomatedTrigger( - autoTriggerType: CodewhispererAutomatedTriggerType, - editor: vscode.TextEditor, - client: DefaultCodeWhispererClient, - config: ConfigurationEntry, - event: vscode.TextDocumentChangeEvent - ): Promise { - if (!editor) { - return - } - - // RecommendationHandler.instance.reportUserDecisionOfRecommendation(editor, -1) - await RecommendationService.instance.generateRecommendation( - client, - editor, - 'AutoTrigger', - config, - autoTriggerType - ) - } -} - -export abstract class DocumentChangedType { - constructor(protected readonly contentChanges: ReadonlyArray) { - this.contentChanges = contentChanges - } - - abstract checkChangeSource(): DocumentChangedSource - - // Enter key should always start with ONE '\n' or '\r\n' and potentially following spaces due to IDE reformat - protected isEnterKey(str: string): boolean { - if (str.length === 0) { - return false - } - return ( - (str.startsWith('\r\n') && str.substring(2).trim() === '') || - (str[0] === '\n' && str.substring(1).trim() === '') - ) - } - - // Tab should consist of space char only ' ' and the length % tabSize should be 0 - protected isTabKey(str: string): boolean { - const tabSize = getTabSizeSetting() - if (str.length % tabSize === 0 && str.trim() === '') { - return true - } - return false - } - - protected isUserTypingSpecialChar(str: string): boolean { - return ['(', '()', '[', '[]', '{', '{}', ':'].includes(str) - } - - protected isSingleLine(str: string): boolean { - let newLineCounts = 0 - for (const ch of str) { - if (ch === '\n') { - newLineCounts += 1 - } - } - - // since pressing Enter key possibly will generate string like '\n ' due to indention - if (this.isEnterKey(str)) { - return true - } - if (newLineCounts >= 1) { - return false - } - return true - } -} - -export class DefaultDocumentChangedType extends DocumentChangedType { - constructor(contentChanges: ReadonlyArray) { - super(contentChanges) - } - - checkChangeSource(): DocumentChangedSource { - if (this.contentChanges.length === 0) { - return DocumentChangedSource.Unknown - } - - // event.contentChanges.length will be 2 when user press Enter key multiple times - if (this.contentChanges.length > 2) { - return DocumentChangedSource.Reformatting - } - - // Case when event.contentChanges.length === 1 - const changedText = this.contentChanges[0].text - - if (this.isSingleLine(changedText)) { - if (changedText === '') { - return DocumentChangedSource.Deletion - } else if (this.isEnterKey(changedText)) { - return DocumentChangedSource.EnterKey - } else if (this.isTabKey(changedText)) { - return DocumentChangedSource.TabKey - } else if (this.isUserTypingSpecialChar(changedText)) { - return DocumentChangedSource.SpecialCharsKey - } else if (changedText.length === 1) { - return DocumentChangedSource.RegularKey - } else if (new RegExp('^[ ]+$').test(changedText)) { - // single line && single place reformat should consist of space chars only - return DocumentChangedSource.Reformatting - } else { - return DocumentChangedSource.Unknown - } - } - - // Won't trigger cwspr on multi-line changes - return DocumentChangedSource.Unknown - } -} - -export enum DocumentChangedSource { - SpecialCharsKey = 'SpecialCharsKey', - RegularKey = 'RegularKey', - TabKey = 'TabKey', - EnterKey = 'EnterKey', - Reformatting = 'Reformatting', - Deletion = 'Deletion', - Unknown = 'Unknown', -} diff --git a/packages/core/src/codewhisperer/service/recommendationHandler.ts b/packages/core/src/codewhisperer/service/recommendationHandler.ts deleted file mode 100644 index 8ab491b32e0..00000000000 --- a/packages/core/src/codewhisperer/service/recommendationHandler.ts +++ /dev/null @@ -1,724 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import { extensionVersion } from '../../shared/vscode/env' -import { RecommendationsList, DefaultCodeWhispererClient, CognitoCredentialsError } from '../client/codewhisperer' -import * as EditorContext from '../util/editorContext' -import * as CodeWhispererConstants from '../models/constants' -import { ConfigurationEntry, GetRecommendationsResponse, vsCodeState } from '../models/model' -import { runtimeLanguageContext } from '../util/runtimeLanguageContext' -import { AWSError } from 'aws-sdk' -import { isAwsError } from '../../shared/errors' -import { TelemetryHelper } from '../util/telemetryHelper' -import { getLogger } from '../../shared/logger/logger' -import { hasVendedIamCredentials } from '../../auth/auth' -import { - asyncCallWithTimeout, - isInlineCompletionEnabled, - isVscHavingRegressionInlineCompletionApi, -} from '../util/commonUtil' -import { showTimedMessage } from '../../shared/utilities/messages' -import { - CodewhispererAutomatedTriggerType, - CodewhispererCompletionType, - CodewhispererGettingStartedTask, - CodewhispererTriggerType, - telemetry, -} from '../../shared/telemetry/telemetry' -import { CodeWhispererCodeCoverageTracker } from '../tracker/codewhispererCodeCoverageTracker' -import { invalidCustomizationMessage } from '../models/constants' -import { getSelectedCustomization, switchToBaseCustomizationAndNotify } from '../util/customizationUtil' -import { session } from '../util/codeWhispererSession' -import { Commands } from '../../shared/vscode/commands2' -import globals from '../../shared/extensionGlobals' -import { noSuggestions, updateInlineLockKey } from '../models/constants' -import AsyncLock from 'async-lock' -import { AuthUtil } from '../util/authUtil' -import { CWInlineCompletionItemProvider } from './inlineCompletionItemProvider' -import { application } from '../util/codeWhispererApplication' -import { openUrl } from '../../shared/utilities/vsCodeUtils' -import { indent } from '../../shared/utilities/textUtilities' -import path from 'path' -import { isIamConnection } from '../../auth/connection' -import { UserWrittenCodeTracker } from '../tracker/userWrittenCodeTracker' - -/** - * This class is for getRecommendation/listRecommendation API calls and its states - * It does not contain UI/UX related logic - */ - -/** - * Commands as a level of indirection so that declare doesn't intercept any registrations for the - * language server implementation. - * - * Otherwise you'll get: - * "Unable to launch amazonq language server: Command "aws.amazonq.rejectCodeSuggestion" has already been declared by the Toolkit" - */ -function createCommands() { - // below commands override VS Code inline completion commands - const prevCommand = Commands.declare('editor.action.inlineSuggest.showPrevious', () => async () => { - await RecommendationHandler.instance.showRecommendation(-1) - }) - const nextCommand = Commands.declare('editor.action.inlineSuggest.showNext', () => async () => { - await RecommendationHandler.instance.showRecommendation(1) - }) - - const rejectCommand = Commands.declare('aws.amazonq.rejectCodeSuggestion', () => async () => { - telemetry.record({ - traceId: TelemetryHelper.instance.traceId, - }) - - await vscode.commands.executeCommand('editor.action.inlineSuggest.hide') - RecommendationHandler.instance.reportUserDecisions(-1) - await Commands.tryExecute('aws.amazonq.refreshAnnotation') - }) - - return { - prevCommand, - nextCommand, - rejectCommand, - } -} - -const lock = new AsyncLock({ maxPending: 1 }) - -export class RecommendationHandler { - public lastInvocationTime: number - // TODO: remove this requestId - public requestId: string - private nextToken: string - private cancellationToken: vscode.CancellationTokenSource - private _onDidReceiveRecommendation: vscode.EventEmitter = new vscode.EventEmitter() - public readonly onDidReceiveRecommendation: vscode.Event = this._onDidReceiveRecommendation.event - private inlineCompletionProvider?: CWInlineCompletionItemProvider - private inlineCompletionProviderDisposable?: vscode.Disposable - private reject: vscode.Disposable - private next: vscode.Disposable - private prev: vscode.Disposable - private _timer?: NodeJS.Timer - documentUri: vscode.Uri | undefined = undefined - - constructor() { - this.requestId = '' - this.nextToken = '' - this.lastInvocationTime = performance.now() - CodeWhispererConstants.invocationTimeIntervalThreshold * 1000 - this.cancellationToken = new vscode.CancellationTokenSource() - this.prev = new vscode.Disposable(() => {}) - this.next = new vscode.Disposable(() => {}) - this.reject = new vscode.Disposable(() => {}) - } - - static #instance: RecommendationHandler - - public static get instance() { - return (this.#instance ??= new this()) - } - - isValidResponse(): boolean { - return session.recommendations.some((r) => r.content.trim() !== '') - } - - async getServerResponse( - triggerType: CodewhispererTriggerType, - isManualTriggerOn: boolean, - promise: Promise - ): Promise { - const timeoutMessage = hasVendedIamCredentials() - ? 'Generate recommendation timeout.' - : 'List recommendation timeout' - if (isManualTriggerOn && triggerType === 'OnDemand' && hasVendedIamCredentials()) { - return vscode.window.withProgress( - { - location: vscode.ProgressLocation.Notification, - title: CodeWhispererConstants.pendingResponse, - cancellable: false, - }, - async () => { - return await asyncCallWithTimeout( - promise, - timeoutMessage, - CodeWhispererConstants.promiseTimeoutLimit * 1000 - ) - } - ) - } - return await asyncCallWithTimeout(promise, timeoutMessage, CodeWhispererConstants.promiseTimeoutLimit * 1000) - } - - async getTaskTypeFromEditorFileName(filePath: string): Promise { - if (filePath.includes('CodeWhisperer_generate_suggestion')) { - return 'autoTrigger' - } else if (filePath.includes('CodeWhisperer_manual_invoke')) { - return 'manualTrigger' - } else if (filePath.includes('CodeWhisperer_use_comments')) { - return 'commentAsPrompt' - } else if (filePath.includes('CodeWhisperer_navigate_suggestions')) { - return 'navigation' - } else if (filePath.includes('Generate_unit_tests')) { - return 'unitTest' - } else { - return undefined - } - } - - async getRecommendations( - client: DefaultCodeWhispererClient, - editor: vscode.TextEditor, - triggerType: CodewhispererTriggerType, - config: ConfigurationEntry, - autoTriggerType?: CodewhispererAutomatedTriggerType, - pagination: boolean = true, - page: number = 0, - generate: boolean = isIamConnection(AuthUtil.instance.conn) - ): Promise { - let invocationResult: 'Succeeded' | 'Failed' = 'Failed' - let errorMessage: string | undefined = undefined - let errorCode: string | undefined = undefined - - if (!editor) { - return Promise.resolve({ - result: invocationResult, - errorMessage: errorMessage, - recommendationCount: 0, - }) - } - let recommendations: RecommendationsList = [] - let requestId = '' - let sessionId = '' - let reason = '' - let startTime = 0 - let latency = 0 - let nextToken = '' - let shouldRecordServiceInvocation = true - session.language = runtimeLanguageContext.getLanguageContext( - editor.document.languageId, - path.extname(editor.document.fileName) - ).language - session.taskType = await this.getTaskTypeFromEditorFileName(editor.document.fileName) - - if (pagination && !generate) { - if (page === 0) { - session.requestContext = await EditorContext.buildListRecommendationRequest( - editor as vscode.TextEditor, - this.nextToken, - config.isSuggestionsWithCodeReferencesEnabled - ) - } else { - session.requestContext = { - request: { - ...session.requestContext.request, - // Putting nextToken assignment in the end so it overwrites the existing nextToken - nextToken: this.nextToken, - }, - supplementalMetadata: session.requestContext.supplementalMetadata, - } - } - } else { - session.requestContext = await EditorContext.buildGenerateRecommendationRequest(editor as vscode.TextEditor) - } - const request = session.requestContext.request - // record preprocessing end time - TelemetryHelper.instance.setPreprocessEndTime() - - // set start pos for non pagination call or first pagination call - if (!pagination || (pagination && page === 0)) { - session.startPos = editor.selection.active - session.startCursorOffset = editor.document.offsetAt(session.startPos) - session.leftContextOfCurrentLine = EditorContext.getLeftContext(editor, session.startPos.line) - session.triggerType = triggerType - session.autoTriggerType = autoTriggerType - - /** - * Validate request - */ - if (!EditorContext.validateRequest(request)) { - getLogger().verbose('Invalid Request: %O', request) - const languageName = request.fileContext.programmingLanguage.languageName - if (!runtimeLanguageContext.isLanguageSupported(languageName)) { - errorMessage = `${languageName} is currently not supported by Amazon Q inline suggestions` - } - return Promise.resolve({ - result: invocationResult, - errorMessage: errorMessage, - recommendationCount: 0, - }) - } - } - - try { - startTime = performance.now() - this.lastInvocationTime = startTime - const mappedReq = runtimeLanguageContext.mapToRuntimeLanguage(request) - const codewhispererPromise = - pagination && !generate - ? client.listRecommendations(mappedReq) - : client.generateRecommendations(mappedReq) - const resp = await this.getServerResponse(triggerType, config.isManualTriggerEnabled, codewhispererPromise) - TelemetryHelper.instance.setSdkApiCallEndTime() - latency = startTime !== 0 ? performance.now() - startTime : 0 - if ('recommendations' in resp) { - recommendations = (resp && resp.recommendations) || [] - } else { - recommendations = (resp && resp.completions) || [] - } - invocationResult = 'Succeeded' - requestId = resp?.$response && resp?.$response?.requestId - nextToken = resp?.nextToken ? resp?.nextToken : '' - sessionId = resp?.$response?.httpResponse?.headers['x-amzn-sessionid'] - TelemetryHelper.instance.setFirstResponseRequestId(requestId) - if (page === 0) { - session.setTimeToFirstRecommendation(performance.now()) - } - if (nextToken === '') { - TelemetryHelper.instance.setAllPaginationEndTime() - } - } catch (error) { - if (error instanceof CognitoCredentialsError) { - shouldRecordServiceInvocation = false - } - if (latency === 0) { - latency = startTime !== 0 ? performance.now() - startTime : 0 - } - getLogger().error('amazonq inline-suggest: Invocation Exception : %s', (error as Error).message) - if (isAwsError(error)) { - errorMessage = error.message - requestId = error.requestId || '' - errorCode = error.code - reason = `CodeWhisperer Invocation Exception: ${error?.code ?? error?.name ?? 'unknown'}` - await this.onThrottlingException(error, triggerType) - - if (error?.code === 'AccessDeniedException' && errorMessage?.includes('no identity-based policy')) { - getLogger().error('amazonq inline-suggest: AccessDeniedException : %s', (error as Error).message) - void vscode.window - .showErrorMessage(`CodeWhisperer: ${error?.message}`, CodeWhispererConstants.settingsLearnMore) - .then(async (resp) => { - if (resp === CodeWhispererConstants.settingsLearnMore) { - void openUrl(vscode.Uri.parse(CodeWhispererConstants.learnMoreUri)) - } - }) - await vscode.commands.executeCommand('aws.amazonq.enableCodeSuggestions', false) - } - } else { - errorMessage = error instanceof Error ? error.message : String(error) - reason = error ? String(error) : 'unknown' - } - } finally { - const timezone = Intl.DateTimeFormat().resolvedOptions().timeZone - - let msg = indent( - `codewhisperer: request-id: ${requestId}, - timestamp(epoch): ${Date.now()}, - timezone: ${timezone}, - datetime: ${new Date().toLocaleString([], { timeZone: timezone })}, - vscode version: '${vscode.version}', - extension version: '${extensionVersion}', - filename: '${EditorContext.getFileName(editor)}', - left context of line: '${session.leftContextOfCurrentLine}', - line number: ${session.startPos.line}, - character location: ${session.startPos.character}, - latency: ${latency} ms. - Recommendations:`, - 4, - true - ).trimStart() - for (const [index, item] of recommendations.entries()) { - msg += `\n ${index.toString().padStart(2, '0')}: ${indent(item.content, 8, true).trim()}` - session.requestIdList.push(requestId) - } - getLogger('nextEditPrediction').debug(`codeWhisper request ${requestId}`) - if (invocationResult === 'Succeeded') { - CodeWhispererCodeCoverageTracker.getTracker(session.language)?.incrementServiceInvocationCount() - UserWrittenCodeTracker.instance.onQFeatureInvoked() - } else { - if ( - (errorMessage?.includes(invalidCustomizationMessage) && errorCode === 'AccessDeniedException') || - errorCode === 'ResourceNotFoundException' - ) { - getLogger() - .debug(`The selected customization is no longer available. Retrying with the default model. - Failed request id: ${requestId}`) - await switchToBaseCustomizationAndNotify() - await this.getRecommendations( - client, - editor, - triggerType, - config, - autoTriggerType, - pagination, - page, - true - ) - } - } - - if (shouldRecordServiceInvocation) { - TelemetryHelper.instance.recordServiceInvocationTelemetry( - requestId, - sessionId, - session.recommendations.length + recommendations.length - 1, - invocationResult, - latency, - session.language, - session.taskType, - reason, - session.requestContext.supplementalMetadata - ) - } - } - - if (this.isCancellationRequested()) { - return Promise.resolve({ - result: invocationResult, - errorMessage: errorMessage, - recommendationCount: session.recommendations.length, - }) - } - - const typedPrefix = editor.document - .getText(new vscode.Range(session.startPos, editor.selection.active)) - .replace('\r\n', '\n') - if (recommendations.length > 0) { - TelemetryHelper.instance.setTypeAheadLength(typedPrefix.length) - // mark suggestions that does not match typeahead when arrival as Discard - // these suggestions can be marked as Showed if typeahead can be removed with new inline API - for (const [i, r] of recommendations.entries()) { - const recommendationIndex = i + session.recommendations.length - if ( - !r.content.startsWith(typedPrefix) && - session.getSuggestionState(recommendationIndex) === undefined - ) { - session.setSuggestionState(recommendationIndex, 'Discard') - } - session.setCompletionType(recommendationIndex, r) - } - session.recommendations = pagination ? session.recommendations.concat(recommendations) : recommendations - if (isInlineCompletionEnabled() && this.hasAtLeastOneValidSuggestion(typedPrefix)) { - this._onDidReceiveRecommendation.fire() - } - } - - this.requestId = requestId - session.sessionId = sessionId - this.nextToken = nextToken - - // send Empty userDecision event if user receives no recommendations in this session at all. - if (invocationResult === 'Succeeded' && nextToken === '') { - // case 1: empty list of suggestion [] - if (session.recommendations.length === 0) { - session.requestIdList.push(requestId) - // Received an empty list of recommendations - TelemetryHelper.instance.recordUserDecisionTelemetryForEmptyList( - session.requestIdList, - sessionId, - page, - runtimeLanguageContext.getLanguageContext( - editor.document.languageId, - path.extname(editor.document.fileName) - ).language, - session.requestContext.supplementalMetadata - ) - } - // case 2: non empty list of suggestion but with (a) empty content or (b) non-matching typeahead - else if (!this.hasAtLeastOneValidSuggestion(typedPrefix)) { - this.reportUserDecisions(-1) - } - } - return Promise.resolve({ - result: invocationResult, - errorMessage: errorMessage, - recommendationCount: session.recommendations.length, - }) - } - - hasAtLeastOneValidSuggestion(typedPrefix: string): boolean { - return session.recommendations.some((r) => r.content.trim() !== '' && r.content.startsWith(typedPrefix)) - } - - cancelPaginatedRequest() { - this.nextToken = '' - this.cancellationToken.cancel() - } - - isCancellationRequested() { - return this.cancellationToken.token.isCancellationRequested - } - - checkAndResetCancellationTokens() { - if (this.isCancellationRequested()) { - this.cancellationToken.dispose() - this.cancellationToken = new vscode.CancellationTokenSource() - this.nextToken = '' - return true - } - return false - } - /** - * Clear recommendation state - */ - clearRecommendations() { - session.requestIdList = [] - session.recommendations = [] - session.suggestionStates = new Map() - session.completionTypes = new Map() - this.requestId = '' - session.sessionId = '' - this.nextToken = '' - session.requestContext.supplementalMetadata = undefined - } - - async clearInlineCompletionStates() { - try { - vsCodeState.isCodeWhispererEditing = false - application()._clearCodeWhispererUIListener.fire() - this.cancelPaginatedRequest() - this.clearRecommendations() - this.disposeInlineCompletion() - await vscode.commands.executeCommand('aws.amazonq.refreshStatusBar') - // fix a regression that requires user to hit Esc twice to clear inline ghost text - // because disposing a provider does not clear the UX - if (isVscHavingRegressionInlineCompletionApi()) { - await vscode.commands.executeCommand('editor.action.inlineSuggest.hide') - } - } finally { - this.clearRejectionTimer() - } - } - - reportDiscardedUserDecisions() { - for (const [i, _] of session.recommendations.entries()) { - session.setSuggestionState(i, 'Discard') - } - this.reportUserDecisions(-1) - } - - /** - * Emits telemetry reflecting user decision for current recommendation. - */ - reportUserDecisions(acceptIndex: number) { - if (session.sessionId === '' || this.requestId === '') { - return - } - TelemetryHelper.instance.recordUserDecisionTelemetry( - session.requestIdList, - session.sessionId, - session.recommendations, - acceptIndex, - session.recommendations.length, - session.completionTypes, - session.suggestionStates, - session.requestContext.supplementalMetadata - ) - if (isInlineCompletionEnabled()) { - this.clearInlineCompletionStates().catch((e) => { - getLogger().error('clearInlineCompletionStates failed: %s', (e as Error).message) - }) - } - } - - hasNextToken(): boolean { - return this.nextToken !== '' - } - - canShowRecommendationInIntelliSense( - editor: vscode.TextEditor, - showPrompt: boolean = false, - response: GetRecommendationsResponse - ): boolean { - const reject = () => { - this.reportUserDecisions(-1) - } - if (!this.isValidResponse()) { - if (showPrompt) { - void showTimedMessage(response.errorMessage ? response.errorMessage : noSuggestions, 3000) - } - reject() - return false - } - // do not show recommendation if cursor is before invocation position - // also mark as Discard - if (editor.selection.active.isBefore(session.startPos)) { - for (const [i, _] of session.recommendations.entries()) { - session.setSuggestionState(i, 'Discard') - } - reject() - return false - } - - // do not show recommendation if typeahead does not match - // also mark as Discard - const typedPrefix = editor.document.getText( - new vscode.Range( - session.startPos.line, - session.startPos.character, - editor.selection.active.line, - editor.selection.active.character - ) - ) - if (!session.recommendations[0].content.startsWith(typedPrefix.trimStart())) { - for (const [i, _] of session.recommendations.entries()) { - session.setSuggestionState(i, 'Discard') - } - reject() - return false - } - return true - } - - async onThrottlingException(awsError: AWSError, triggerType: CodewhispererTriggerType) { - if ( - awsError.code === 'ThrottlingException' && - awsError.message.includes(CodeWhispererConstants.throttlingMessage) - ) { - if (triggerType === 'OnDemand') { - void vscode.window.showErrorMessage(CodeWhispererConstants.freeTierLimitReached) - } - vsCodeState.isFreeTierLimitReached = true - } - } - - public disposeInlineCompletion() { - this.inlineCompletionProviderDisposable?.dispose() - this.inlineCompletionProvider = undefined - } - - private disposeCommandOverrides() { - this.prev.dispose() - this.reject.dispose() - this.next.dispose() - } - - // These commands override the vs code inline completion commands - // They are subscribed when suggestion starts and disposed when suggestion is accepted/rejected - // to avoid impacting other plugins or user who uses this API - private registerCommandOverrides() { - const { prevCommand, nextCommand, rejectCommand } = createCommands() - this.prev = prevCommand.register() - this.next = nextCommand.register() - this.reject = rejectCommand.register() - } - - subscribeSuggestionCommands() { - this.disposeCommandOverrides() - this.registerCommandOverrides() - globals.context.subscriptions.push(this.prev) - globals.context.subscriptions.push(this.next) - globals.context.subscriptions.push(this.reject) - } - - async showRecommendation(indexShift: number, noSuggestionVisible: boolean = false) { - await lock.acquire(updateInlineLockKey, async () => { - if (!vscode.window.state.focused) { - this.reportDiscardedUserDecisions() - return - } - const inlineCompletionProvider = new CWInlineCompletionItemProvider( - this.inlineCompletionProvider?.getActiveItemIndex, - indexShift, - session.recommendations, - this.requestId, - session.startPos, - this.nextToken - ) - this.inlineCompletionProviderDisposable?.dispose() - // when suggestion is active, registering a new provider will let VS Code invoke inline API automatically - this.inlineCompletionProviderDisposable = vscode.languages.registerInlineCompletionItemProvider( - Object.assign([], CodeWhispererConstants.platformLanguageIds), - inlineCompletionProvider - ) - this.inlineCompletionProvider = inlineCompletionProvider - - if (isVscHavingRegressionInlineCompletionApi() && !noSuggestionVisible) { - // fix a regression in new VS Code when disposing and re-registering - // a new provider does not auto refresh the inline suggestion widget - // by manually refresh it - await vscode.commands.executeCommand('editor.action.inlineSuggest.hide') - await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger') - } - if (noSuggestionVisible) { - await vscode.commands.executeCommand(`editor.action.inlineSuggest.trigger`) - this.sendPerceivedLatencyTelemetry() - } - }) - } - - async onEditorChange() { - this.reportUserDecisions(-1) - } - - async onFocusChange() { - this.reportUserDecisions(-1) - } - - async onCursorChange(e: vscode.TextEditorSelectionChangeEvent) { - // we do not want to reset the states for keyboard events because they can be typeahead - if ( - e.kind !== vscode.TextEditorSelectionChangeKind.Keyboard && - vscode.window.activeTextEditor === e.textEditor - ) { - application()._clearCodeWhispererUIListener.fire() - // when cursor change due to mouse movement we need to reset the active item index for inline - if (e.kind === vscode.TextEditorSelectionChangeKind.Mouse) { - this.inlineCompletionProvider?.clearActiveItemIndex() - } - } - } - - isSuggestionVisible(): boolean { - return this.inlineCompletionProvider?.getActiveItemIndex !== undefined - } - - async tryShowRecommendation() { - const editor = vscode.window.activeTextEditor - if (editor === undefined) { - return - } - if (this.isSuggestionVisible()) { - // do not force refresh the tooltip to avoid suggestion "flashing" - return - } - if ( - editor.selection.active.isBefore(session.startPos) || - editor.document.uri.fsPath !== this.documentUri?.fsPath - ) { - for (const [i, _] of session.recommendations.entries()) { - session.setSuggestionState(i, 'Discard') - } - this.reportUserDecisions(-1) - } else if (session.recommendations.length > 0) { - await this.showRecommendation(0, true) - } - } - - private clearRejectionTimer() { - if (this._timer !== undefined) { - clearInterval(this._timer) - this._timer = undefined - } - } - - private sendPerceivedLatencyTelemetry() { - if (vscode.window.activeTextEditor) { - const languageContext = runtimeLanguageContext.getLanguageContext( - vscode.window.activeTextEditor.document.languageId, - vscode.window.activeTextEditor.document.fileName.substring( - vscode.window.activeTextEditor.document.fileName.lastIndexOf('.') + 1 - ) - ) - telemetry.codewhisperer_perceivedLatency.emit({ - codewhispererRequestId: this.requestId, - codewhispererSessionId: session.sessionId, - codewhispererTriggerType: session.triggerType, - codewhispererCompletionType: session.getCompletionType(0), - codewhispererCustomizationArn: getSelectedCustomization().arn, - codewhispererLanguage: languageContext.language, - duration: performance.now() - this.lastInvocationTime, - passive: true, - credentialStartUrl: AuthUtil.instance.startUrl, - result: 'Succeeded', - }) - } - } -} diff --git a/packages/core/src/codewhisperer/service/recommendationService.ts b/packages/core/src/codewhisperer/service/recommendationService.ts deleted file mode 100644 index de78b435913..00000000000 --- a/packages/core/src/codewhisperer/service/recommendationService.ts +++ /dev/null @@ -1,122 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -import * as vscode from 'vscode' -import { ConfigurationEntry, GetRecommendationsResponse } from '../models/model' -import { isInlineCompletionEnabled } from '../util/commonUtil' -import { - CodewhispererAutomatedTriggerType, - CodewhispererTriggerType, - telemetry, -} from '../../shared/telemetry/telemetry' -import { InlineCompletionService } from '../service/inlineCompletionService' -import { ClassifierTrigger } from './classifierTrigger' -import { DefaultCodeWhispererClient } from '../client/codewhisperer' -import { randomUUID } from '../../shared/crypto' -import { TelemetryHelper } from '../util/telemetryHelper' -import { AuthUtil } from '../util/authUtil' - -export interface SuggestionActionEvent { - readonly editor: vscode.TextEditor | undefined - readonly isRunning: boolean - readonly triggerType: CodewhispererTriggerType - readonly response: GetRecommendationsResponse | undefined -} - -export class RecommendationService { - static #instance: RecommendationService - - private _isRunning: boolean = false - get isRunning() { - return this._isRunning - } - - private _onSuggestionActionEvent = new vscode.EventEmitter() - get suggestionActionEvent(): vscode.Event { - return this._onSuggestionActionEvent.event - } - - private _acceptedSuggestionCount: number = 0 - get acceptedSuggestionCount() { - return this._acceptedSuggestionCount - } - - private _totalValidTriggerCount: number = 0 - get totalValidTriggerCount() { - return this._totalValidTriggerCount - } - - public static get instance() { - return (this.#instance ??= new RecommendationService()) - } - - incrementAcceptedCount() { - this._acceptedSuggestionCount++ - } - - incrementValidTriggerCount() { - this._totalValidTriggerCount++ - } - - async generateRecommendation( - client: DefaultCodeWhispererClient, - editor: vscode.TextEditor, - triggerType: CodewhispererTriggerType, - config: ConfigurationEntry, - autoTriggerType?: CodewhispererAutomatedTriggerType, - event?: vscode.TextDocumentChangeEvent - ) { - // TODO: should move all downstream auth check(inlineCompletionService, recommendationHandler etc) to here(upstream) instead of spreading everywhere - if (AuthUtil.instance.isConnected() && AuthUtil.instance.requireProfileSelection()) { - return - } - - if (this._isRunning) { - return - } - - /** - * Use an existing trace ID if invoked through a command (e.g., manual invocation), - * otherwise generate a new trace ID - */ - const traceId = telemetry.attributes?.traceId ?? randomUUID() - TelemetryHelper.instance.setTraceId(traceId) - await telemetry.withTraceId(async () => { - if (isInlineCompletionEnabled()) { - if (triggerType === 'OnDemand') { - ClassifierTrigger.instance.recordClassifierResultForManualTrigger(editor) - } - - this._isRunning = true - let response: GetRecommendationsResponse | undefined = undefined - - try { - this._onSuggestionActionEvent.fire({ - editor: editor, - isRunning: true, - triggerType: triggerType, - response: undefined, - }) - - response = await InlineCompletionService.instance.getPaginatedRecommendation( - client, - editor, - triggerType, - config, - autoTriggerType, - event - ) - } finally { - this._isRunning = false - this._onSuggestionActionEvent.fire({ - editor: editor, - isRunning: false, - triggerType: triggerType, - response: response, - }) - } - } - }, traceId) - } -} diff --git a/packages/core/src/codewhisperer/service/referenceLogViewProvider.ts b/packages/core/src/codewhisperer/service/referenceLogViewProvider.ts index 9ec20b8cb44..2f6e9231697 100644 --- a/packages/core/src/codewhisperer/service/referenceLogViewProvider.ts +++ b/packages/core/src/codewhisperer/service/referenceLogViewProvider.ts @@ -4,13 +4,15 @@ */ import * as vscode from 'vscode' -import { References } from '../client/codewhisperer' import { LicenseUtil } from '../util/licenseUtil' import * as CodeWhispererConstants from '../models/constants' import { CodeWhispererSettings } from '../util/codewhispererSettings' import globals from '../../shared/extensionGlobals' import { AuthUtil } from '../util/authUtil' import { session } from '../util/codeWhispererSession' +import CodeWhispererClient from '../client/codewhispererclient' +import CodeWhispererUserClient from '../client/codewhispereruserclient' +import { InlineCompletionItemWithReferences } from '@aws/language-server-runtimes-types' export class ReferenceLogViewProvider implements vscode.WebviewViewProvider { public static readonly viewType = 'aws.codeWhisperer.referenceLog' @@ -52,28 +54,23 @@ export class ReferenceLogViewProvider implements vscode.WebviewViewProvider { } } - public static getReferenceLog(recommendation: string, references: References, editor: vscode.TextEditor): string { + public static getReferenceLog(recommendation: string, references: Reference[], editor: vscode.TextEditor): string { const filePath = editor.document.uri.path const time = new Date().toLocaleString() let text = `` for (const reference of references) { + const standardReference = toStandardReference(reference) if ( - reference.recommendationContentSpan === undefined || - reference.recommendationContentSpan.start === undefined || - reference.recommendationContentSpan.end === undefined + standardReference.position === undefined || + standardReference.position.start === undefined || + standardReference.position.end === undefined ) { continue } - const code = recommendation.substring( - reference.recommendationContentSpan.start, - reference.recommendationContentSpan.end - ) - const firstCharLineNumber = - editor.document.positionAt(session.startCursorOffset + reference.recommendationContentSpan.start).line + - 1 - const lastCharLineNumber = - editor.document.positionAt(session.startCursorOffset + reference.recommendationContentSpan.end - 1) - .line + 1 + const { start, end } = standardReference.position + const code = recommendation.substring(start, end) + const firstCharLineNumber = editor.document.positionAt(session.startCursorOffset + start).line + 1 + const lastCharLineNumber = editor.document.positionAt(session.startCursorOffset + end - 1).line + 1 let lineInfo = `` if (firstCharLineNumber === lastCharLineNumber) { lineInfo = `(line at ${firstCharLineNumber})` @@ -84,11 +81,11 @@ export class ReferenceLogViewProvider implements vscode.WebviewViewProvider { text += `And ` } - let license = `${reference.licenseName}` - let repository = reference.repository?.length ? reference.repository : 'unknown' - if (reference.url?.length) { - repository = `${reference.repository}` - license = `${reference.licenseName || 'unknown'}` + let license = `${standardReference.licenseName}` + let repository = standardReference.repository?.length ? standardReference.repository : 'unknown' + if (standardReference.url?.length) { + repository = `${standardReference.repository}` + license = `${standardReference.licenseName || 'unknown'}` } text += @@ -119,7 +116,7 @@ export class ReferenceLogViewProvider implements vscode.WebviewViewProvider { let prompt = '' if (showPrompt) { - if (AuthUtil.instance.isEnterpriseSsoInUse()) { + if (AuthUtil.instance.isIdcConnection()) { prompt = CodeWhispererConstants.referenceLogPromptTextEnterpriseSSO } else { prompt = CodeWhispererConstants.referenceLogPromptText @@ -144,3 +141,48 @@ export class ReferenceLogViewProvider implements vscode.WebviewViewProvider { ` } } + +/** + * Reference log needs to support references directly from CW, as well as those from Flare. These references have different shapes, so we standarize them here. + */ +type GetInnerType = T extends (infer U)[] ? U : never +type Reference = + | CodeWhispererClient.Reference + | CodeWhispererUserClient.Reference + | GetInnerType + +type StandardizedReference = { + licenseName?: string + position?: { + start?: number + end?: number + } + repository?: string + url?: string +} + +/** + * Convert a general reference to the standardized format expected by the reference log. + * @param ref + * @returns + */ +function toStandardReference(ref: Reference): StandardizedReference { + const isCWReference = (ref: any) => ref.recommendationContentSpan !== undefined + + if (isCWReference(ref)) { + const castRef = ref as CodeWhispererClient.Reference + return { + licenseName: castRef.licenseName!, + position: { start: castRef.recommendationContentSpan?.start, end: castRef.recommendationContentSpan?.end }, + repository: castRef.repository, + url: castRef.url, + } + } + const castRef = ref as GetInnerType + return { + licenseName: castRef.licenseName, + position: { start: castRef.position?.startCharacter, end: castRef.position?.endCharacter }, + repository: castRef.referenceName, + url: castRef.referenceUrl, + } +} diff --git a/packages/core/src/codewhisperer/service/securityIssueHoverProvider.ts b/packages/core/src/codewhisperer/service/securityIssueHoverProvider.ts index b82c10063e6..7934ff8f26c 100644 --- a/packages/core/src/codewhisperer/service/securityIssueHoverProvider.ts +++ b/packages/core/src/codewhisperer/service/securityIssueHoverProvider.ts @@ -44,7 +44,7 @@ export class SecurityIssueHoverProvider implements vscode.HoverProvider { detectorId: issue.detectorId, ruleId: issue.ruleId, includesFix: !!issue.suggestedFixes.length, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, autoDetected: issue.autoDetected, }) TelemetryHelper.instance.sendCodeScanRemediationsEvent( diff --git a/packages/core/src/codewhisperer/service/securityScanHandler.ts b/packages/core/src/codewhisperer/service/securityScanHandler.ts index b83fdbebb1a..5195e8a3e4c 100644 --- a/packages/core/src/codewhisperer/service/securityScanHandler.ts +++ b/packages/core/src/codewhisperer/service/securityScanHandler.ts @@ -300,7 +300,7 @@ export async function getPresignedUrlAndUpload( span.record({ amazonqUploadIntent: uploadIntent, amazonqRepositorySize: zipMetadata.srcPayloadSizeInBytes, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, }) const srcReq: CreateUploadUrlRequest = { contentMd5: getMd5(zipMetadata.zipFilePath), diff --git a/packages/core/src/codewhisperer/service/statusBar.ts b/packages/core/src/codewhisperer/service/statusBar.ts new file mode 100644 index 00000000000..5e10df164f3 --- /dev/null +++ b/packages/core/src/codewhisperer/service/statusBar.ts @@ -0,0 +1,147 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +import * as vscode from 'vscode' +import { CodeSuggestionsState } from '../models/model' +import { AuthUtil } from '../util/authUtil' +import { getSelectedCustomization } from '../util/customizationUtil' +import { codicon, getIcon } from '../../shared/icons' +import { Commands } from '../../shared/vscode/commands2' +import { listCodeWhispererCommandsId } from '../ui/statusBarMenu' + +export class CodeWhispererStatusBarManager { + private statusBar: CodeWhispererStatusBar + + constructor(statusBar: CodeWhispererStatusBar = CodeWhispererStatusBar.instance) { + this.statusBar = statusBar + + CodeSuggestionsState.instance.onDidChangeState(() => { + return this.refreshStatusBar() + }) + } + + static #instance: CodeWhispererStatusBarManager + + public static get instance() { + return (this.#instance ??= new this()) + } + + /** Updates the status bar to represent the latest CW state */ + refreshStatusBar() { + if (AuthUtil.instance.isConnected()) { + if (AuthUtil.instance.regionProfileManager.requireProfileSelection()) { + return this.setState('needsProfile') + } + return this.setState('ok') + } else if (AuthUtil.instance.isConnectionExpired()) { + return this.setState('expired') + } else { + return this.setState('notConnected') + } + } + + /** + * Sets the status bar in to a "loading state", effectively showing + * the spinning circle. + * + * When loading is done, call {@link refreshStatusBar} to update the + * status bar to the latest state. + */ + async setLoading(): Promise { + await this.setState('loading') + } + + private async setState(state: keyof typeof states) { + switch (state) { + case 'loading': { + await this.statusBar.setState('loading') + break + } + case 'ok': { + await this.statusBar.setState('ok', CodeSuggestionsState.instance.isSuggestionsEnabled()) + break + } + case 'expired': { + await this.statusBar.setState('expired') + break + } + case 'notConnected': { + await this.statusBar.setState('notConnected') + break + } + case 'needsProfile': { + await this.statusBar.setState('needsProfile') + break + } + } + } +} + +/** The states that the completion service can be in */ +const states = { + loading: 'loading', + ok: 'ok', + expired: 'expired', + notConnected: 'notConnected', + needsProfile: 'needsProfile', +} as const + +class CodeWhispererStatusBar { + protected statusBar: vscode.StatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 1) + + static #instance: CodeWhispererStatusBar + static get instance() { + return (this.#instance ??= new this()) + } + + protected constructor() {} + + async setState(state: keyof Omit): Promise + async setState(status: keyof Pick, isSuggestionsEnabled: boolean): Promise + async setState(status: keyof typeof states, isSuggestionsEnabled?: boolean): Promise { + const statusBar = this.statusBar + statusBar.command = listCodeWhispererCommandsId + statusBar.backgroundColor = undefined + + const title = 'Amazon Q' + switch (status) { + case 'loading': { + const selectedCustomization = getSelectedCustomization() + statusBar.text = codicon` ${getIcon('vscode-loading~spin')} ${title}${ + selectedCustomization.arn === '' ? '' : ` | ${selectedCustomization.name}` + }` + break + } + case 'ok': { + const selectedCustomization = getSelectedCustomization() + const icon = isSuggestionsEnabled ? getIcon('vscode-debug-start') : getIcon('vscode-debug-pause') + statusBar.text = codicon`${icon} ${title}${ + selectedCustomization.arn === '' ? '' : ` | ${selectedCustomization.name}` + }` + break + } + + case 'expired': { + statusBar.text = codicon` ${getIcon('vscode-debug-disconnect')} ${title}` + statusBar.backgroundColor = new vscode.ThemeColor('statusBarItem.warningBackground') + break + } + case 'needsProfile': + case 'notConnected': + statusBar.text = codicon` ${getIcon('vscode-chrome-close')} ${title}` + statusBar.backgroundColor = new vscode.ThemeColor('statusBarItem.errorBackground') + break + } + + statusBar.show() + } +} + +/** In this module due to circular dependency issues */ +export const refreshStatusBar = Commands.declare( + { id: 'aws.amazonq.refreshStatusBar', logging: false }, + () => async () => { + await CodeWhispererStatusBarManager.instance.refreshStatusBar() + } +) diff --git a/packages/core/src/codewhisperer/tracker/codewhispererCodeCoverageTracker.ts b/packages/core/src/codewhisperer/tracker/codewhispererCodeCoverageTracker.ts deleted file mode 100644 index 0989f022245..00000000000 --- a/packages/core/src/codewhisperer/tracker/codewhispererCodeCoverageTracker.ts +++ /dev/null @@ -1,319 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import { getLogger } from '../../shared/logger/logger' -import * as CodeWhispererConstants from '../models/constants' -import globals from '../../shared/extensionGlobals' -import { vsCodeState } from '../models/model' -import { CodewhispererLanguage, telemetry } from '../../shared/telemetry/telemetry' -import { runtimeLanguageContext } from '../util/runtimeLanguageContext' -import { TelemetryHelper } from '../util/telemetryHelper' -import { AuthUtil } from '../util/authUtil' -import { getSelectedCustomization } from '../util/customizationUtil' -import { codeWhispererClient as client } from '../client/codewhisperer' -import { isAwsError } from '../../shared/errors' -import { getUnmodifiedAcceptedTokens } from '../util/commonUtil' - -interface CodeWhispererToken { - range: vscode.Range - text: string - accepted: number -} - -const autoClosingKeystrokeInputs = ['[]', '{}', '()', '""', "''"] - -/** - * This singleton class is mainly used for calculating the code written by codeWhisperer - * TODO: Remove this tracker, uses user written code tracker instead. - * This is kept in codebase for server side backward compatibility until service fully switch to user written code - */ -export class CodeWhispererCodeCoverageTracker { - private _acceptedTokens: { [key: string]: CodeWhispererToken[] } - private _totalTokens: { [key: string]: number } - private _timer?: NodeJS.Timer - private _startTime: number - private _language: CodewhispererLanguage - private _serviceInvocationCount: number - - private constructor(language: CodewhispererLanguage) { - this._acceptedTokens = {} - this._totalTokens = {} - this._startTime = 0 - this._language = language - this._serviceInvocationCount = 0 - } - - public get serviceInvocationCount(): number { - return this._serviceInvocationCount - } - - public get acceptedTokens(): { [key: string]: CodeWhispererToken[] } { - return this._acceptedTokens - } - - public get totalTokens(): { [key: string]: number } { - return this._totalTokens - } - - public isActive(): boolean { - return TelemetryHelper.instance.isTelemetryEnabled() && AuthUtil.instance.isConnected() - } - - public incrementServiceInvocationCount() { - this._serviceInvocationCount += 1 - } - - public flush() { - if (!this.isActive()) { - this._totalTokens = {} - this._acceptedTokens = {} - this.closeTimer() - return - } - try { - this.emitCodeWhispererCodeContribution() - } catch (error) { - getLogger().error(`Encountered ${error} when emitting code contribution metric`) - } - } - - // TODO: Improve the range tracking of the accepted recommendation - // TODO: use the editor of the filename, not the current editor - public updateAcceptedTokensCount(editor: vscode.TextEditor) { - const filename = editor.document.fileName - if (filename in this._acceptedTokens) { - for (let i = 0; i < this._acceptedTokens[filename].length; i++) { - const oldText = this._acceptedTokens[filename][i].text - const newText = editor.document.getText(this._acceptedTokens[filename][i].range) - this._acceptedTokens[filename][i].accepted = getUnmodifiedAcceptedTokens(oldText, newText) - } - } - } - - public emitCodeWhispererCodeContribution() { - let totalTokens = 0 - for (const filename in this._totalTokens) { - totalTokens += this._totalTokens[filename] - } - if (vscode.window.activeTextEditor) { - this.updateAcceptedTokensCount(vscode.window.activeTextEditor) - } - // the accepted characters without counting user modification - let acceptedTokens = 0 - // the accepted characters after calculating user modification - let unmodifiedAcceptedTokens = 0 - for (const filename in this._acceptedTokens) { - for (const v of this._acceptedTokens[filename]) { - if (filename in this._totalTokens && this._totalTokens[filename] >= v.accepted) { - unmodifiedAcceptedTokens += v.accepted - acceptedTokens += v.text.length - } - } - } - const percentCount = ((acceptedTokens / totalTokens) * 100).toFixed(2) - const percentage = Math.round(parseInt(percentCount)) - const selectedCustomization = getSelectedCustomization() - if (this._serviceInvocationCount <= 0) { - getLogger().debug(`Skip emiting code contribution metric`) - return - } - telemetry.codewhisperer_codePercentage.emit({ - codewhispererTotalTokens: totalTokens, - codewhispererLanguage: this._language, - codewhispererAcceptedTokens: unmodifiedAcceptedTokens, - codewhispererSuggestedTokens: acceptedTokens, - codewhispererPercentage: percentage ? percentage : 0, - successCount: this._serviceInvocationCount, - codewhispererCustomizationArn: selectedCustomization.arn === '' ? undefined : selectedCustomization.arn, - credentialStartUrl: AuthUtil.instance.startUrl, - }) - - client - .sendTelemetryEvent({ - telemetryEvent: { - codeCoverageEvent: { - customizationArn: selectedCustomization.arn === '' ? undefined : selectedCustomization.arn, - programmingLanguage: { - languageName: runtimeLanguageContext.toRuntimeLanguage(this._language), - }, - acceptedCharacterCount: acceptedTokens, - unmodifiedAcceptedCharacterCount: unmodifiedAcceptedTokens, - totalCharacterCount: totalTokens, - timestamp: new Date(Date.now()), - }, - }, - profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, - }) - .then() - .catch((error) => { - let requestId: string | undefined - if (isAwsError(error)) { - requestId = error.requestId - } - - getLogger().debug( - `Failed to sendTelemetryEvent to CodeWhisperer, requestId: ${requestId ?? ''}, message: ${ - error.message - }` - ) - }) - } - - private tryStartTimer() { - if (this._timer !== undefined) { - return - } - const currentDate = new globals.clock.Date() - this._startTime = currentDate.getTime() - this._timer = setTimeout(() => { - try { - const currentTime = new globals.clock.Date().getTime() - const delay: number = CodeWhispererConstants.defaultCheckPeriodMillis - const diffTime: number = this._startTime + delay - if (diffTime <= currentTime) { - let totalTokens = 0 - for (const filename in this._totalTokens) { - totalTokens += this._totalTokens[filename] - } - if (totalTokens > 0) { - this.flush() - } else { - getLogger().debug( - `CodeWhispererCodeCoverageTracker: skipped telemetry due to empty tokens array` - ) - } - } - } catch (e) { - getLogger().verbose(`Exception Thrown from CodeWhispererCodeCoverageTracker: ${e}`) - } finally { - this.resetTracker() - this.closeTimer() - } - }, CodeWhispererConstants.defaultCheckPeriodMillis) - } - - private resetTracker() { - this._totalTokens = {} - this._acceptedTokens = {} - this._startTime = 0 - this._serviceInvocationCount = 0 - } - - private closeTimer() { - if (this._timer !== undefined) { - clearTimeout(this._timer) - this._timer = undefined - } - } - - public addAcceptedTokens(filename: string, token: CodeWhispererToken) { - if (!(filename in this._acceptedTokens)) { - this._acceptedTokens[filename] = [] - } - this._acceptedTokens[filename].push(token) - } - - public addTotalTokens(filename: string, count: number) { - if (!(filename in this._totalTokens)) { - this._totalTokens[filename] = 0 - } - this._totalTokens[filename] += count - if (this._totalTokens[filename] < 0) { - this._totalTokens[filename] = 0 - } - } - - public countAcceptedTokens(range: vscode.Range, text: string, filename: string) { - if (!this.isActive()) { - return - } - // generate accepted recommendation token and stored in collection - this.addAcceptedTokens(filename, { range: range, text: text, accepted: text.length }) - this.addTotalTokens(filename, text.length) - } - - // For below 2 edge cases - // 1. newline character with indentation - // 2. 2 character insertion of closing brackets - public getCharacterCountFromComplexEvent(e: vscode.TextDocumentChangeEvent) { - function countChanges(cond: boolean, text: string): number { - if (!cond) { - return 0 - } - if ((text.startsWith('\n') || text.startsWith('\r\n')) && text.trim().length === 0) { - return 1 - } - if (autoClosingKeystrokeInputs.includes(text)) { - return 2 - } - return 0 - } - if (e.contentChanges.length === 2) { - const text1 = e.contentChanges[0].text - const text2 = e.contentChanges[1].text - const text2Count = countChanges(text1.length === 0, text2) - const text1Count = countChanges(text2.length === 0, text1) - return text2Count > 0 ? text2Count : text1Count - } else if (e.contentChanges.length === 1) { - return countChanges(true, e.contentChanges[0].text) - } - return 0 - } - - public isFromUserKeystroke(e: vscode.TextDocumentChangeEvent) { - return e.contentChanges.length === 1 && e.contentChanges[0].text.length === 1 - } - - public countTotalTokens(e: vscode.TextDocumentChangeEvent) { - // ignore no contentChanges. ignore contentChanges from other plugins (formatters) - // only include contentChanges from user keystroke input(one character input). - // Also ignore deletion events due to a known issue of tracking deleted CodeWhiperer tokens. - if (!runtimeLanguageContext.isLanguageSupported(e.document.languageId) || vsCodeState.isCodeWhispererEditing) { - return - } - // a user keystroke input can be - // 1. content change with 1 character insertion - // 2. newline character with indentation - // 3. 2 character insertion of closing brackets - if (this.isFromUserKeystroke(e)) { - this.tryStartTimer() - this.addTotalTokens(e.document.fileName, 1) - } else if (this.getCharacterCountFromComplexEvent(e) !== 0) { - this.tryStartTimer() - const characterIncrease = this.getCharacterCountFromComplexEvent(e) - this.addTotalTokens(e.document.fileName, characterIncrease) - } - // also include multi character input within 50 characters (not from CWSPR) - else if ( - e.contentChanges.length === 1 && - e.contentChanges[0].text.length > 1 && - TelemetryHelper.instance.lastSuggestionInDisplay !== e.contentChanges[0].text - ) { - const multiCharInputSize = e.contentChanges[0].text.length - - // select 50 as the cut-off threshold for counting user input. - // ignore all white space multi char input, this usually comes from reformat. - if (multiCharInputSize < 50 && e.contentChanges[0].text.trim().length > 0) { - this.addTotalTokens(e.document.fileName, multiCharInputSize) - } - } - } - - public static readonly instances = new Map() - - public static getTracker(language: string): CodeWhispererCodeCoverageTracker | undefined { - if (!runtimeLanguageContext.isLanguageSupported(language)) { - return undefined - } - const cwsprLanguage = runtimeLanguageContext.normalizeLanguage(language) - if (!cwsprLanguage) { - return undefined - } - const instance = this.instances.get(cwsprLanguage) ?? new this(cwsprLanguage) - this.instances.set(cwsprLanguage, instance) - return instance - } -} diff --git a/packages/core/src/codewhisperer/tracker/codewhispererTracker.ts b/packages/core/src/codewhisperer/tracker/codewhispererTracker.ts index ca19c87505f..c20dfb3c900 100644 --- a/packages/core/src/codewhisperer/tracker/codewhispererTracker.ts +++ b/packages/core/src/codewhisperer/tracker/codewhispererTracker.ts @@ -111,7 +111,7 @@ export class CodeWhispererTracker { cwsprChatConversationId: suggestion.conversationID, cwsprChatMessageId: suggestion.messageID, cwsprChatModificationPercentage: percentage ? percentage : 0, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, } telemetry.amazonq_modifyCode.emit(event) @@ -138,7 +138,7 @@ export class CodeWhispererTracker { codewhispererModificationPercentage: percentage ? percentage : 0, codewhispererCompletionType: suggestion.completionType, codewhispererLanguage: suggestion.language, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, codewhispererCharactersAccepted: suggestion.originalString.length, codewhispererCharactersModified: 0, // TODO: currently we don't have an accurate number for this field with existing implementation }) diff --git a/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts index c3e46bdc78e..1b887e587d5 100644 --- a/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts +++ b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts @@ -28,6 +28,7 @@ import { AuthUtil } from '../util/authUtil' import { submitFeedback } from '../../feedback/vue/submitFeedback' import { focusAmazonQPanel } from '../../codewhispererChat/commands/registerCommands' import { isWeb } from '../../shared/extensionGlobals' +import { builderIdRegion, builderIdStartUrl } from '../../auth/sso/constants' export function createAutoSuggestions(running: boolean): DataQuickPickItem<'autoSuggestions'> { const labelResume = localize('AWS.codewhisperer.resumeCodeWhispererNode.label', 'Resume Auto-Suggestions') @@ -178,7 +179,7 @@ export function createGettingStarted(): DataQuickPickItem<'gettingStarted'> { export function createSignout(): DataQuickPickItem<'signout'> { const label = localize('AWS.codewhisperer.signoutNode.label', 'Sign Out') const icon = getIcon('vscode-export') - const connection = AuthUtil.instance.isBuilderIdInUse() ? 'AWS Builder ID' : 'IAM Identity Center' + const connection = AuthUtil.instance.isBuilderIdConnection() ? 'AWS Builder ID' : 'IAM Identity Center' return { data: 'signout', @@ -252,7 +253,7 @@ export function createSignIn(): DataQuickPickItem<'signIn'> { if (isWeb()) { // TODO: nkomonen, call a Command instead onClick = () => { - void AuthUtil.instance.connectToAwsBuilderId() + void AuthUtil.instance.login(builderIdStartUrl, builderIdRegion) } } diff --git a/packages/core/src/codewhisperer/ui/statusBarMenu.ts b/packages/core/src/codewhisperer/ui/statusBarMenu.ts index 2ad14a81df0..d8c05270073 100644 --- a/packages/core/src/codewhisperer/ui/statusBarMenu.ts +++ b/packages/core/src/codewhisperer/ui/statusBarMenu.ts @@ -43,7 +43,7 @@ function getAmazonQCodeWhispererNodes() { return [createSignIn(), createLearnMore()] } - if (AuthUtil.instance.isConnected() && AuthUtil.instance.requireProfileSelection()) { + if (AuthUtil.instance.isConnected() && AuthUtil.instance.regionProfileManager.requireProfileSelection()) { return [] } @@ -72,12 +72,12 @@ function getAmazonQCodeWhispererNodes() { // Security scans createSeparator('Code Reviews'), - ...(AuthUtil.instance.isBuilderIdInUse() ? [] : [createAutoScans(autoScansEnabled)]), + ...(AuthUtil.instance.isBuilderIdConnection() ? [] : [createAutoScans(autoScansEnabled)]), createSecurityScan(), // Amazon Q + others createSeparator('Other Features'), - ...(AuthUtil.instance.isValidEnterpriseSsoInUse() && AuthUtil.instance.isCustomizationFeatureEnabled + ...(AuthUtil.instance.isIdcConnection() && AuthUtil.instance.isCustomizationFeatureEnabled ? [createSelectCustomization()] : []), switchToAmazonQNode(), @@ -85,7 +85,7 @@ function getAmazonQCodeWhispererNodes() { } export function getQuickPickItems(): DataQuickPickItem[] { - const isUsingEnterpriseSso = AuthUtil.instance.isValidEnterpriseSsoInUse() + const isUsingEnterpriseSso = AuthUtil.instance.isIdcConnection() const regionProfile = AuthUtil.instance.regionProfileManager.activeRegionProfile const children = [ @@ -104,7 +104,7 @@ export function getQuickPickItems(): DataQuickPickItem[] { // Add settings and signout createSeparator(), createSettingsNode(), - ...(isUsingEnterpriseSso && regionProfile ? [createSelectRegionProfileNode(regionProfile)] : []), + ...(AuthUtil.instance.isIdcConnection() && regionProfile ? [createSelectRegionProfileNode(regionProfile)] : []), ...(AuthUtil.instance.isConnected() && !hasVendedIamCredentials() && !hasVendedCredentialsFromMetadata() ? [createSignout()] : []), diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 10acbe16424..1419eaa4772 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -5,360 +5,208 @@ import * as vscode from 'vscode' import * as localizedText from '../../shared/localizedText' -import { Auth } from '../../auth/auth' -import { ToolkitError, isNetworkError, tryRun } from '../../shared/errors' -import { getSecondaryAuth, setScopes } from '../../auth/secondaryAuth' -import { isSageMaker } from '../../shared/extensionUtilities' +import * as nls from 'vscode-nls' +import { fs } from '../../shared/fs/fs' +import * as path from 'path' +import { ToolkitError } from '../../shared/errors' import { AmazonQPromptSettings } from '../../shared/settings' import { scopesCodeWhispererCore, - createBuilderIdProfile, - hasScopes, - SsoConnection, - createSsoProfile, - Connection, - isIamConnection, - isSsoConnection, - isBuilderIdConnection, scopesCodeWhispererChat, scopesFeatureDev, scopesGumby, - isIdcSsoConnection, + TelemetryMetadata, + scopesSsoAccountAccess, + hasScopes, + SsoProfile, + StoredProfile, hasExactScopes, - getTelemetryMetadataForConn, - ProfileNotFoundError, } from '../../auth/connection' import { getLogger } from '../../shared/logger/logger' -import { Commands, placeholder } from '../../shared/vscode/commands2' +import { Commands } from '../../shared/vscode/commands2' import { vsCodeState } from '../models/model' -import { onceChanged, once } from '../../shared/utilities/functionUtils' -import { indent } from '../../shared/utilities/textUtilities' import { showReauthenticateMessage } from '../../shared/utilities/messages' import { showAmazonQWalkthroughOnce } from '../../amazonq/onboardingPage/walkthrough' import { setContext } from '../../shared/vscode/setContext' -import { isInDevEnv } from '../../shared/vscode/env' import { openUrl } from '../../shared/utilities/vsCodeUtils' -import * as nls from 'vscode-nls' -const localize = nls.loadMessageBundle() import { telemetry } from '../../shared/telemetry/telemetry' -import { asStringifiedStack } from '../../shared/telemetry/spans' -import { withTelemetryContext } from '../../shared/telemetry/util' -import { focusAmazonQPanel } from '../../codewhispererChat/commands/registerCommands' -import { throttle } from 'lodash' +import { AuthStateEvent, cacheChangedEvent, LanguageClientAuth, LoginTypes, SsoLogin } from '../../auth/auth2' +import { builderIdStartUrl, internalStartUrl } from '../../auth/sso/constants' +import { VSCODE_EXTENSION_ID } from '../../shared/extensions' import { RegionProfileManager } from '../region/regionProfileManager' +import { AuthFormId } from '../../login/webview/vue/types' +import { getEnvironmentSpecificMemento } from '../../shared/utilities/mementos' +import { getCacheDir, getFlareCacheFileName, getRegistrationCacheFile, getTokenCacheFile } from '../../auth/sso/cache' +import { notifySelectDeveloperProfile } from '../region/utils' +import { once } from '../../shared/utilities/functionUtils' +import { CancellationTokenSource, SsoTokenSourceKind } from '@aws/language-server-runtimes/server-interface' + +const localize = nls.loadMessageBundle() /** Backwards compatibility for connections w pre-chat scopes */ export const codeWhispererCoreScopes = [...scopesCodeWhispererCore] export const codeWhispererChatScopes = [...codeWhispererCoreScopes, ...scopesCodeWhispererChat] export const amazonQScopes = [...codeWhispererChatScopes, ...scopesGumby, ...scopesFeatureDev] -/** - * "Core" are the CW scopes that existed before the addition of new scopes - * for Amazon Q. - */ -export const isValidCodeWhispererCoreConnection = (conn?: Connection): conn is Connection => { - return ( - (isSageMaker() && isIamConnection(conn)) || (isSsoConnection(conn) && hasScopes(conn, codeWhispererCoreScopes)) - ) -} -/** Superset that includes all of CodeWhisperer + Amazon Q */ -export const isValidAmazonQConnection = (conn?: Connection): conn is Connection => { - return ( - (isSageMaker() && isIamConnection(conn)) || - ((isSsoConnection(conn) || isBuilderIdConnection(conn)) && - isValidCodeWhispererCoreConnection(conn) && - hasScopes(conn, amazonQScopes)) - ) +/** AuthProvider interface for the auth functionality needed by RegionProfileManager */ +export interface IAuthProvider { + isConnected(): boolean + isBuilderIdConnection(): boolean + isIdcConnection(): boolean + isSsoSession(): boolean + getToken(): Promise + readonly profileName: string + readonly connection?: { region: string; startUrl: string } } -const authClassName = 'AuthQ' - -export class AuthUtil { - static #instance: AuthUtil - protected static readonly logIfChanged = onceChanged((s: string) => getLogger().info(s)) +/** + * Handles authentication within Amazon Q. + * Amazon Q only supports a single connection at a time. + */ +export class AuthUtil implements IAuthProvider { + public readonly profileName = VSCODE_EXTENSION_ID.amazonq + protected logger = getLogger('amazonqAuth') - private reauthenticatePromptShown: boolean = false - private _isCustomizationFeatureEnabled: boolean = false + public readonly regionProfileManager: RegionProfileManager - // user should only see that screen once. - // TODO: move to memento - public hasAlreadySeenMigrationAuthScreen: boolean = false + // IAM login currently not supported + private session: SsoLogin - public get isCustomizationFeatureEnabled(): boolean { - return this._isCustomizationFeatureEnabled + static create(lspAuth: LanguageClientAuth) { + return (this.#instance ??= new this(lspAuth)) } - // This boolean controls whether the Select Customization node will be visible. A change to this value - // means that the old UX was wrong and must refresh the devTool tree. - public set isCustomizationFeatureEnabled(value: boolean) { - if (this._isCustomizationFeatureEnabled === value) { - return + static #instance: AuthUtil + public static get instance() { + if (!this.#instance) { + throw new ToolkitError('AuthUtil not ready. Was it initialized with a running LSP?') } - this._isCustomizationFeatureEnabled = value - void Commands.tryExecute('aws.amazonq.refreshStatusBar') + return this.#instance } - public readonly secondaryAuth = getSecondaryAuth( - this.auth, - 'codewhisperer', - 'Amazon Q', - isValidCodeWhispererCoreConnection - ) - public readonly restore = () => this.secondaryAuth.restoreConnection() - - public constructor( - public readonly auth = Auth.instance, - public readonly regionProfileManager = new RegionProfileManager(() => this.conn) - ) {} - - public initCodeWhispererHooks = once(() => { - this.auth.onDidChangeConnectionState(async (e) => { - getLogger().info(`codewhisperer: connection changed to ${e.state}: ${e.id}`) - if (e.state !== 'authenticating') { - await this.refreshCodeWhisperer() - } - - await this.setVscodeContextProps() - }) - - this.secondaryAuth.onDidChangeActiveConnection(async () => { - getLogger().info(`codewhisperer: active connection changed`) - if (this.isValidEnterpriseSsoInUse()) { - void vscode.commands.executeCommand('aws.amazonq.notifyNewCustomizations') - await this.regionProfileManager.restoreProfileSelection() - } - vsCodeState.isFreeTierLimitReached = false - await Promise.all([ - // onDidChangeActiveConnection may trigger before these modules are activated. - Commands.tryExecute('aws.amazonq.refreshStatusBar'), - Commands.tryExecute('aws.amazonq.updateReferenceLog'), - ]) - - await this.setVscodeContextProps() - - // To check valid connection - if (this.isValidEnterpriseSsoInUse() || (this.isBuilderIdInUse() && !this.isConnectionExpired())) { - await showAmazonQWalkthroughOnce() - } - - if (!this.isConnected()) { - await this.regionProfileManager.invalidateProfile(this.regionProfileManager.activeRegionProfile?.arn) - await this.regionProfileManager.clearCache() - } - }) + private constructor(private readonly lspAuth: LanguageClientAuth) { + this.session = new SsoLogin(this.profileName, this.lspAuth) + this.onDidChangeConnectionState((e: AuthStateEvent) => this.stateChangeHandler(e)) + this.regionProfileManager = new RegionProfileManager(this) this.regionProfileManager.onDidChangeRegionProfile(async () => { await this.setVscodeContextProps() }) - }) - - public async setVscodeContextProps() { - // if users are "pending profile selection", they're not fully connected and require profile selection for Q usage - // requireProfileSelection() always returns false for builderID users - await setContext('aws.codewhisperer.connected', this.isConnected() && !this.requireProfileSelection()) - const doShowAmazonQLoginView = - !this.isConnected() || this.isConnectionExpired() || this.requireProfileSelection() - await setContext('aws.amazonq.showLoginView', doShowAmazonQLoginView) - await setContext('aws.codewhisperer.connectionExpired', this.isConnectionExpired()) - await setContext('aws.amazonq.connectedSsoIdc', isIdcSsoConnection(this.conn)) - } - - public reformatStartUrl(startUrl: string | undefined) { - return !startUrl ? undefined : startUrl.replace(/[\/#]+$/g, '') - } - - // current active cwspr connection - public get conn() { - return this.secondaryAuth.activeConnection + lspAuth.registerCacheWatcher(async (event: cacheChangedEvent) => await this.cacheChangedHandler(event)) } - // TODO: move this to the shared auth.ts - public get startUrl(): string | undefined { - // Reformat the url to remove any trailing '/' and `#` - // e.g. https://view.awsapps.com/start/# will become https://view.awsapps.com/start - return isSsoConnection(this.conn) ? this.reformatStartUrl(this.conn?.startUrl) : undefined + // Do NOT use this in production code, only used for testing + static destroy(): void { + this.#instance = undefined as any } - public get isUsingSavedConnection() { - return this.conn !== undefined && this.secondaryAuth.hasSavedConnection + isSsoSession() { + return this.session.loginType === LoginTypes.SSO } - public isConnected(): boolean { - return this.conn !== undefined - } - - public isEnterpriseSsoInUse(): boolean { - const conn = this.conn - // we have an sso that isn't builder id, must be IdC by process of elimination - const isUsingEnterpriseSso = conn?.type === 'sso' && !isBuilderIdConnection(conn) - return conn !== undefined && isUsingEnterpriseSso + /** + * HACK: Ideally we'd put {@link notifySelectDeveloperProfile} in to {@link restore}. + * But because {@link refreshState} is only called if !isConnected, we cannot do it since + * {@link notifySelectDeveloperProfile} needs {@link refreshState} to run so it can set + * the Bearer Token in the LSP first. + */ + didStartSignedIn = false + + async restore() { + await this.session.restore() + this.didStartSignedIn = this.isConnected() + + // HACK: We noticed that if calling `refreshState()` here when the user was already signed in, something broke. + // So as a solution we only call it if they were not already signed in. + // + // But in the case where a user was already signed in, we allow `session.restore()` to trigger `refreshState()` through + // event emitters. + // This is unoptimal since `refreshState()` should be able to be called multiple times and still work. + // + // Because of this edge case, when `restore()` is called we cannot assume all Auth is setup when this function returns, + // since we may still be waiting on the event emitter to trigger the expected functions. + // + // TODO: Figure out why removing the if statement below causes things to break. Maybe we just need to + // promisify the call and any subsequent callers will not make a redundant call. + if (!this.didStartSignedIn) { + await this.refreshState() + } } - // If there is an active SSO connection - public isValidEnterpriseSsoInUse(): boolean { - return this.isEnterpriseSsoInUse() && !this.isConnectionExpired() - } + async login(startUrl: string, region: string) { + const response = await this.session.login({ startUrl, region, scopes: amazonQScopes }) + await showAmazonQWalkthroughOnce() - public isBuilderIdInUse(): boolean { - return this.conn !== undefined && isBuilderIdConnection(this.conn) + return response } - @withTelemetryContext({ name: 'connectToAwsBuilderId', class: authClassName }) - public async connectToAwsBuilderId(): Promise { - let conn = (await this.auth.listConnections()).find(isBuilderIdConnection) - - if (!conn) { - conn = await this.auth.createConnection(createBuilderIdProfile(amazonQScopes)) - } else if (!isValidAmazonQConnection(conn)) { - conn = await this.secondaryAuth.addScopes(conn, amazonQScopes) + reauthenticate() { + if (!this.isSsoSession()) { + throw new ToolkitError('Cannot reauthenticate non-SSO session.') } - if (this.auth.getConnectionState(conn) === 'invalid') { - conn = await this.auth.reauthenticate(conn) - } - - return (await this.secondaryAuth.useNewConnection(conn)) as SsoConnection + return this.session.reauthenticate() } - @withTelemetryContext({ name: 'connectToEnterpriseSso', class: authClassName }) - public async connectToEnterpriseSso(startUrl: string, region: string): Promise { - let conn = (await this.auth.listConnections()).find( - (conn): conn is SsoConnection => - isSsoConnection(conn) && conn.startUrl.toLowerCase() === startUrl.toLowerCase() - ) - - if (!conn) { - conn = await this.auth.createConnection(createSsoProfile(startUrl, region, amazonQScopes)) - } else if (!isValidAmazonQConnection(conn)) { - conn = await this.secondaryAuth.addScopes(conn, amazonQScopes) - } - - if (this.auth.getConnectionState(conn) === 'invalid') { - conn = await this.auth.reauthenticate(conn) + logout() { + if (!this.isSsoSession()) { + // Only SSO requires logout + return } - - return (await this.secondaryAuth.useNewConnection(conn)) as SsoConnection + this.lspAuth.deleteBearerToken() + return this.session.logout() } - public static get instance() { - if (this.#instance !== undefined) { - return this.#instance + async getToken() { + if (this.isSsoSession()) { + return (await this.session.getToken()).token + } else { + throw new ToolkitError('Cannot get token for non-SSO session.') } - - const self = (this.#instance = new this()) - return self } - @withTelemetryContext({ name: 'getBearerToken', class: authClassName }) - public async getBearerToken(): Promise { - await this.restore() - - if (this.conn === undefined) { - throw new ToolkitError('No connection found', { code: 'NoConnection' }) - } - - if (!isSsoConnection(this.conn)) { - throw new ToolkitError('Connection is not an SSO connection', { code: 'BadConnectionType' }) - } - - try { - const bearerToken = await this.conn.getToken() - return bearerToken.accessToken - } catch (err) { - if (err instanceof ProfileNotFoundError) { - // Expected that connection would be deleted by conn.getToken() - void focusAmazonQPanel.execute(placeholder, 'profileNotFoundSignout') - } - throw err - } + get connection() { + return this.session.data } - @withTelemetryContext({ name: 'getCredentials', class: authClassName }) - public async getCredentials() { - await this.restore() - - if (this.conn === undefined) { - throw new ToolkitError('No connection found', { code: 'NoConnection' }) - } - - if (!isIamConnection(this.conn)) { - throw new ToolkitError('Connection is not an IAM connection', { code: 'BadConnectionType' }) - } - - return this.conn.getCredentials() + getAuthState() { + return this.session.getConnectionState() } - public isConnectionValid(log: boolean = true): boolean { - const connectionValid = this.conn !== undefined && !this.secondaryAuth.isConnectionExpired - - if (log) { - this.logConnection() - } - - return connectionValid + isConnected() { + return this.getAuthState() === 'connected' } - public isConnectionExpired(log: boolean = true): boolean { - const connectionExpired = - this.secondaryAuth.isConnectionExpired && - this.conn !== undefined && - isValidCodeWhispererCoreConnection(this.conn) - - if (log) { - this.logConnection() - } - - return connectionExpired + isConnectionExpired() { + return this.getAuthState() === 'expired' } - requireProfileSelection(): boolean { - if (isBuilderIdConnection(this.conn)) { - return false - } - return isIdcSsoConnection(this.conn) && this.regionProfileManager.activeRegionProfile === undefined + isBuilderIdConnection() { + return this.connection?.startUrl === builderIdStartUrl } - private logConnection() { - const logStr = indent( - `codewhisperer: connection states - connection isValid=${this.isConnectionValid(false)}, - connection isValidCodewhispererCoreConnection=${isValidCodeWhispererCoreConnection(this.conn)}, - connection isExpired=${this.isConnectionExpired(false)}, - secondaryAuth isExpired=${this.secondaryAuth.isConnectionExpired}, - connection isUndefined=${this.conn === undefined}`, - 4, - true - ) - - AuthUtil.logIfChanged(logStr) + isIdcConnection() { + return Boolean(this.connection?.startUrl && this.connection?.startUrl !== builderIdStartUrl) } - @withTelemetryContext({ name: 'reauthenticate', class: authClassName }) - public async reauthenticate() { - try { - if (this.conn?.type !== 'sso') { - return - } - - if (!hasExactScopes(this.conn, amazonQScopes)) { - const conn = await setScopes(this.conn, amazonQScopes, this.auth) - await this.secondaryAuth.useNewConnection(conn) - } + isInternalAmazonUser(): boolean { + return this.isConnected() && this.connection?.startUrl === internalStartUrl + } - await this.auth.reauthenticate(this.conn) - } catch (err) { - throw ToolkitError.chain(err, 'Unable to authenticate connection') - } finally { - await this.setVscodeContextProps() - } + onDidChangeConnectionState(handler: (e: AuthStateEvent) => any) { + return this.session.onDidChangeConnectionState(handler) } - public async refreshCodeWhisperer() { - vsCodeState.isFreeTierLimitReached = false - await Commands.tryExecute('aws.amazonq.refreshStatusBar') + public async setVscodeContextProps(state = this.getAuthState()) { + await setContext('aws.codewhisperer.connected', state === 'connected') + const showAmazonQLoginView = + !this.isConnected() || this.isConnectionExpired() || this.regionProfileManager.requireProfileSelection() + await setContext('aws.amazonq.showLoginView', showAmazonQLoginView) + await setContext('aws.amazonq.connectedSsoIdc', this.isIdcConnection()) + await setContext('aws.codewhisperer.connectionExpired', state === 'expired') } - @withTelemetryContext({ name: 'showReauthenticatePrompt', class: authClassName }) + private reauthenticatePromptShown: boolean = false public async showReauthenticatePrompt(isAutoTrigger?: boolean) { if (isAutoTrigger && this.reauthenticatePromptShown) { return @@ -379,6 +227,26 @@ export class AuthUtil { } } + private _isCustomizationFeatureEnabled: boolean = false + public get isCustomizationFeatureEnabled(): boolean { + return this._isCustomizationFeatureEnabled + } + + // This boolean controls whether the Select Customization node will be visible. A change to this value + // means that the old UX was wrong and must refresh the devTool tree. + public set isCustomizationFeatureEnabled(value: boolean) { + if (this._isCustomizationFeatureEnabled === value) { + return + } + this._isCustomizationFeatureEnabled = value + void Commands.tryExecute('aws.amazonq.refreshStatusBar') + } + + public async notifyReauthenticate(isAutoTrigger?: boolean) { + void this.showReauthenticatePrompt(isAutoTrigger) + await this.setVscodeContextProps() + } + public async notifySessionConfiguration() { const suppressId = 'amazonQSessionConfigurationMessage' const settings = AmazonQPromptSettings.instance @@ -411,175 +279,203 @@ export class AuthUtil { }) } - @withTelemetryContext({ name: 'notifyReauthenticate', class: authClassName }) - public async notifyReauthenticate(isAutoTrigger?: boolean) { - void this.showReauthenticatePrompt(isAutoTrigger) - await this.setVscodeContextProps() + private async cacheChangedHandler(event: cacheChangedEvent) { + this.logger.debug(`Cache change event received: ${event}`) + if (event === 'delete') { + await this.logout() + } else if (event === 'create') { + await this.restore() + } } - public isValidCodeTransformationAuthUser(): boolean { - return (this.isEnterpriseSsoInUse() || this.isBuilderIdInUse()) && this.isConnectionValid() + private async stateChangeHandler(e: AuthStateEvent) { + if (e.state === 'refreshed') { + const params = this.isSsoSession() ? (await this.session.getToken()).updateCredentialsParams : undefined + await this.lspAuth.updateBearerToken(params!) + return + } else { + this.logger.info(`codewhisperer: connection changed to ${e.state}`) + await this.refreshState(e.state) + } } - /** - * Asynchronously returns a snapshot of the overall auth state of CodeWhisperer + Chat features. - * It guarantees the latest state is correct at the risk of modifying connection state. - * If this guarantee is not required, use sync method getChatAuthStateSync() - * - * By default, network errors are ignored when determining auth state since they may be silently - * recoverable later. - * - * THROTTLE: This function is called in rapid succession by Amazon Q features and can lead to - * a barrage of disk access and/or token refreshes. We throttle to deal with this. - * - * Note we do an explicit cast of the return type due to Lodash types incorrectly indicating - * a FeatureAuthState or undefined can be returned. But since we set `leading: true` - * it will always return FeatureAuthState - */ - public getChatAuthState = throttle(() => this._getChatAuthState(), 2000, { - leading: true, - }) as () => Promise - /** - * IMPORTANT: Only use this if you do NOT want to swallow network errors, otherwise use {@link getChatAuthState()} - * @param ignoreNetErr swallows network errors - */ - @withTelemetryContext({ name: 'getChatAuthState', class: authClassName }) - public async _getChatAuthState(ignoreNetErr: boolean = true): Promise { - // The state of the connection may not have been properly validated - // and the current state we see may be stale, so refresh for latest state. - if (ignoreNetErr) { - await tryRun( - () => this.auth.refreshConnectionState(this.conn), - (err) => !isNetworkError(err), - 'getChatAuthState: Cannot refresh connection state due to network error: %s' - ) - } else { - await this.auth.refreshConnectionState(this.conn) + private async refreshState(state = this.getAuthState()) { + if (state === 'expired' || state === 'notConnected') { + this.lspAuth.deleteBearerToken() + if (this.isIdcConnection()) { + await this.regionProfileManager.invalidateProfile(this.regionProfileManager.activeRegionProfile?.arn) + await this.regionProfileManager.clearCache() + } + } + if (state === 'connected') { + const bearerTokenParams = (await this.session.getToken()).updateCredentialsParams + await this.lspAuth.updateBearerToken(bearerTokenParams) + + if (this.isIdcConnection()) { + await this.regionProfileManager.restoreProfileSelection() + } } - return this.getChatAuthStateSync(this.conn) + // regardless of state, send message at startup if user needs to select a Developer Profile + void this.tryNotifySelectDeveloperProfile() + + vsCodeState.isFreeTierLimitReached = false + await this.setVscodeContextProps(state) + await Promise.all([ + Commands.tryExecute('aws.amazonq.refreshStatusBar'), + Commands.tryExecute('aws.amazonq.updateReferenceLog'), + ]) + + if (state === 'connected' && this.isIdcConnection()) { + void vscode.commands.executeCommand('aws.amazonq.notifyNewCustomizations') + } } - /** - * Synchronously returns a snapshot of the overall auth state of CodeWhisperer + Chat features without - * validating or modifying the connection state. It is possible that the connection - * is invalid/valid, but the current state displays something else. To guarantee the true state, - * use async method getChatAuthState() - */ - public getChatAuthStateSync(conn = this.conn): FeatureAuthState { - if (conn === undefined) { - return buildFeatureAuthState(AuthStates.disconnected) + private tryNotifySelectDeveloperProfile = once(async () => { + if (this.regionProfileManager.requireProfileSelection() && this.didStartSignedIn) { + await notifySelectDeveloperProfile() } + }) - if (!isSsoConnection(conn) && !isSageMaker()) { - throw new ToolkitError(`Connection "${conn.id}" is not a valid type: ${conn.type}`) + async getTelemetryMetadata(): Promise { + if (!this.isConnected()) { + return { + id: 'undefined', + } } - // default to expired to indicate reauth is needed if unmodified - const state: FeatureAuthState = buildFeatureAuthState(AuthStates.expired) + if (this.isSsoSession()) { + const ssoSessionDetails = (await this.session.getProfile()).ssoSession?.settings + return { + authScopes: ssoSessionDetails?.sso_registration_scopes?.join(','), + credentialSourceId: AuthUtil.instance.isBuilderIdConnection() ? 'awsId' : 'iamIdentityCenter', + credentialStartUrl: AuthUtil.instance.connection?.startUrl, + awsRegion: AuthUtil.instance.connection?.region, + } + } else if (!AuthUtil.instance.isSsoSession) { + return { + credentialSourceId: 'sharedCredentials', + } + } + + throw new Error('getTelemetryMetadataForConn() called with unknown connection type') + } - if (this.isConnectionExpired()) { - return state + async getAuthFormIds(): Promise { + if (!this.isConnected()) { + return [] } - if (isBuilderIdConnection(conn) || isIdcSsoConnection(conn) || isSageMaker()) { - // TODO: refactor - if (isValidCodeWhispererCoreConnection(conn)) { - if (this.requireProfileSelection()) { - state[Features.codewhispererCore] = AuthStates.pendingProfileSelection - } else { - state[Features.codewhispererCore] = AuthStates.connected - } - } - if (isValidAmazonQConnection(conn)) { - if (this.requireProfileSelection()) { - for (const v of Object.values(Features)) { - state[v as Feature] = AuthStates.pendingProfileSelection - } - } else { - for (const v of Object.values(Features)) { - state[v as Feature] = AuthStates.connected - } - } + const authIds: AuthFormId[] = [] + let connType: 'builderId' | 'identityCenter' + + // TODO: update when there is IAM support + if (!this.isSsoSession()) { + return ['credentials'] + } else if (this.isBuilderIdConnection()) { + connType = 'builderId' + } else if (this.isIdcConnection()) { + connType = 'identityCenter' + const ssoSessionDetails = (await this.session.getProfile()).ssoSession?.settings + if (hasScopes(ssoSessionDetails?.sso_registration_scopes ?? [], scopesSsoAccountAccess)) { + authIds.push('identityCenterExplorer') } + } else { + return ['unknown'] } + authIds.push(`${connType}CodeWhisperer`) - return state + return authIds } /** - * Edge Case: Due to a change in behaviour/functionality, there are potential extra - * auth connections that the Amazon Q extension has cached. We need to remove these - * as they are irrelevant to the Q extension and can cause issues. + * Migrates existing SSO connections to the LSP identity server by updating the cache files + * + * @param clientName - The client name to use for the new registration cache file + * @returns A Promise that resolves when the migration is complete + * @throws Error if file operations fail during migration */ - public async clearExtraConnections(): Promise { - const currentQConn = this.conn - // Q currently only maintains 1 connection at a time, so we assume everything else is extra. - // IMPORTANT: In the case Q starts to manage multiple connections, this implementation will need to be updated. - const allOtherConnections = (await this.auth.listConnections()).filter((c) => c.id !== currentQConn?.id) - for (const conn of allOtherConnections) { - getLogger().warn(`forgetting extra amazon q connection: %O`, conn) - await telemetry.auth_modifyConnection.run( - async () => { - telemetry.record({ - connectionState: Auth.instance.getConnectionState(conn) ?? 'undefined', - source: asStringifiedStack(telemetry.getFunctionStack()), - ...(await getTelemetryMetadataForConn(conn)), - }) + async migrateSsoConnectionToLsp(clientName: string) { + const memento = getEnvironmentSpecificMemento() + const key = 'auth.profiles' + const profiles: { readonly [id: string]: StoredProfile } | undefined = memento.get(key) - if (isInDevEnv()) { - telemetry.record({ action: 'forget' }) - // in a Dev Env the connection may be used by code catalyst, so we forget instead of fully deleting - await this.auth.forgetConnection(conn) - } else { - telemetry.record({ action: 'delete' }) - await this.auth.deleteConnection(conn) - } + let toImport: SsoProfile | undefined + let profileId: string | undefined + + if (!profiles) { + return + } + + try { + // Try go get token from LSP auth. If available, skip migration and delete old auth profile + const token = await this.lspAuth.getSsoToken( + { + kind: SsoTokenSourceKind.IamIdentityCenter, + profileName: this.profileName, }, - { functionId: { name: 'clearExtraConnections', class: authClassName } } + false, + new CancellationTokenSource().token ) + if (token) { + this.logger.info('existing LSP auth connection found. Skipping migration') + await memento.update(key, undefined) + return + } + } catch { + this.logger.info('unable to get token from LSP auth, proceeding migration') } - } -} -export type FeatureAuthState = { [feature in Feature]: AuthState } -export type Feature = (typeof Features)[keyof typeof Features] -export type AuthState = (typeof AuthStates)[keyof typeof AuthStates] + this.logger.info('checking for old SSO connections') + for (const [id, p] of Object.entries(profiles)) { + if (p.type === 'sso' && hasExactScopes(p.scopes ?? [], amazonQScopes)) { + toImport = p + profileId = id + if (p.metadata.connectionState === 'valid') { + break + } + } + } -export const AuthStates = { - /** The current connection is working and supports this feature. */ - connected: 'connected', - /** No connection exists, so this feature cannot be used*/ - disconnected: 'disconnected', - /** - * The current connection exists, but needs to be reauthenticated for this feature to work - * - * Look to use {@link AuthUtil.reauthenticate()} - */ - expired: 'expired', - /** - * A connection exists, but does not support this feature. - * - * Eg: We are currently using Builder ID, but must use Identity Center. - */ - unsupported: 'unsupported', - /** - * The current connection exists and isn't expired, - * but fetching/refreshing the token resulted in a network error. - */ - connectedWithNetworkError: 'connectedWithNetworkError', - pendingProfileSelection: 'pendingProfileSelection', -} as const -const Features = { - codewhispererCore: 'codewhispererCore', - codewhispererChat: 'codewhispererChat', - amazonQ: 'amazonQ', -} as const - -function buildFeatureAuthState(state: AuthState): FeatureAuthState { - return { - codewhispererCore: state, - codewhispererChat: state, - amazonQ: state, + if (toImport && profileId) { + this.logger.info('migrating SSO connection to LSP identity server...') + + const registrationKey = { + startUrl: toImport.startUrl, + region: toImport.ssoRegion, + scopes: amazonQScopes, + } + + await this.session.updateProfile(registrationKey) + + const cacheDir = getCacheDir() + + const fromRegistrationFile = getRegistrationCacheFile(cacheDir, registrationKey) + const toRegistrationFile = path.join( + cacheDir, + getFlareCacheFileName( + JSON.stringify({ + region: toImport.ssoRegion, + startUrl: toImport.startUrl, + tool: clientName, + }) + ) + ) + + const fromTokenFile = getTokenCacheFile(cacheDir, profileId) + const toTokenFile = path.join(cacheDir, getFlareCacheFileName(this.profileName)) + + try { + await fs.rename(fromRegistrationFile, toRegistrationFile) + await fs.rename(fromTokenFile, toTokenFile) + this.logger.debug('Successfully renamed registration and token files') + } catch (err) { + this.logger.error(`Failed to rename files during migration: ${err}`) + throw err + } + + this.logger.info('successfully migrated SSO connection to LSP identity server') + await memento.update(key, undefined) + } } } diff --git a/packages/core/src/codewhisperer/util/closingBracketUtil.ts b/packages/core/src/codewhisperer/util/closingBracketUtil.ts deleted file mode 100644 index 466ca31a0b9..00000000000 --- a/packages/core/src/codewhisperer/util/closingBracketUtil.ts +++ /dev/null @@ -1,262 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import * as CodeWhispererConstants from '../models/constants' - -interface bracketMapType { - [k: string]: string -} - -const quotes = ["'", '"', '`'] -const parenthesis = ['(', '[', '{', ')', ']', '}', '<', '>'] - -const closeToOpen: bracketMapType = { - ')': '(', - ']': '[', - '}': '{', - '>': '<', -} - -const openToClose: bracketMapType = { - '(': ')', - '[': ']', - '{': '}', - '<': '>', -} - -/** - * LeftContext | Recommendation | RightContext - * This function aims to resolve symbols which are redundant and need to be removed - * The high level logic is as followed - * 1. Pair non-paired closing symbols(parenthesis, brackets, quotes) existing in the "recommendation" with non-paired symbols existing in the "leftContext" - * 2. Remove non-paired closing symbols existing in the "rightContext" - * @param endPosition: end position of the effective recommendation written by CodeWhisperer - * @param startPosition: start position of the effective recommendation by CodeWhisperer - * - * for example given file context ('|' is where we trigger the service): - * anArray.pu| - * recommendation returned: "sh(element);" - * typeahead: "sh(" - * the effective recommendation written by CodeWhisperer: "element);" - */ -export async function handleExtraBrackets( - editor: vscode.TextEditor, - endPosition: vscode.Position, - startPosition: vscode.Position -) { - const recommendation = editor.document.getText(new vscode.Range(startPosition, endPosition)) - const endOffset = editor.document.offsetAt(endPosition) - const startOffset = editor.document.offsetAt(startPosition) - const leftContext = editor.document.getText( - new vscode.Range( - startPosition, - editor.document.positionAt(Math.max(startOffset - CodeWhispererConstants.charactersLimit, 0)) - ) - ) - - const rightContext = editor.document.getText( - new vscode.Range( - editor.document.positionAt(endOffset), - editor.document.positionAt(endOffset + CodeWhispererConstants.charactersLimit) - ) - ) - const bracketsToRemove = getBracketsToRemove( - editor, - recommendation, - leftContext, - rightContext, - endPosition, - startPosition - ) - - const quotesToRemove = getQuotesToRemove( - editor, - recommendation, - leftContext, - rightContext, - endPosition, - startPosition - ) - - const symbolsToRemove = [...bracketsToRemove, ...quotesToRemove] - - if (symbolsToRemove.length) { - await removeBracketsFromRightContext(editor, symbolsToRemove, endPosition) - } -} - -const removeBracketsFromRightContext = async ( - editor: vscode.TextEditor, - idxToRemove: number[], - endPosition: vscode.Position -) => { - const offset = editor.document.offsetAt(endPosition) - - await editor.edit( - (editBuilder) => { - for (const idx of idxToRemove) { - const range = new vscode.Range( - editor.document.positionAt(offset + idx), - editor.document.positionAt(offset + idx + 1) - ) - editBuilder.delete(range) - } - }, - { undoStopAfter: false, undoStopBefore: false } - ) -} - -function getBracketsToRemove( - editor: vscode.TextEditor, - recommendation: string, - leftContext: string, - rightContext: string, - end: vscode.Position, - start: vscode.Position -) { - const unpairedClosingsInReco = nonClosedClosingParen(recommendation) - const unpairedOpeningsInLeftContext = nonClosedOpneingParen(leftContext, unpairedClosingsInReco.length) - const unpairedClosingsInRightContext = nonClosedClosingParen(rightContext) - - const toRemove: number[] = [] - - let i = 0 - let j = 0 - let k = 0 - while (i < unpairedOpeningsInLeftContext.length && j < unpairedClosingsInReco.length) { - const opening = unpairedOpeningsInLeftContext[i] - const closing = unpairedClosingsInReco[j] - - const isPaired = closeToOpen[closing.char] === opening.char - const rightContextCharToDelete = unpairedClosingsInRightContext[k] - - if (isPaired) { - if (rightContextCharToDelete && rightContextCharToDelete.char === closing.char) { - const rightContextStart = editor.document.offsetAt(end) + 1 - const symbolPosition = editor.document.positionAt( - rightContextStart + rightContextCharToDelete.strOffset - ) - const lineCnt = recommendation.split('\n').length - 1 - const isSameline = symbolPosition.line - lineCnt === start.line - - if (isSameline) { - toRemove.push(rightContextCharToDelete.strOffset) - } - - k++ - } - } - - i++ - j++ - } - - return toRemove -} - -function getQuotesToRemove( - editor: vscode.TextEditor, - recommendation: string, - leftContext: string, - rightContext: string, - endPosition: vscode.Position, - startPosition: vscode.Position -) { - let leftQuote: string | undefined = undefined - let leftIndex: number | undefined = undefined - for (let i = leftContext.length - 1; i >= 0; i--) { - const char = leftContext[i] - if (quotes.includes(char)) { - leftQuote = char - leftIndex = leftContext.length - i - break - } - } - - let rightQuote: string | undefined = undefined - let rightIndex: number | undefined = undefined - for (let i = 0; i < rightContext.length; i++) { - const char = rightContext[i] - if (quotes.includes(char)) { - rightQuote = char - rightIndex = i - break - } - } - - let quoteCountInReco = 0 - if (leftQuote && rightQuote && leftQuote === rightQuote) { - for (const char of recommendation) { - if (quotes.includes(char) && char === leftQuote) { - quoteCountInReco++ - } - } - } - - if (leftIndex !== undefined && rightIndex !== undefined && quoteCountInReco % 2 !== 0) { - const p = editor.document.positionAt(editor.document.offsetAt(endPosition) + rightIndex) - - if (endPosition.line === startPosition.line && endPosition.line === p.line) { - return [rightIndex] - } - } - - return [] -} - -function nonClosedOpneingParen(str: string, cnt?: number): { char: string; strOffset: number }[] { - const resultSet: { char: string; strOffset: number }[] = [] - const stack: string[] = [] - - for (let i = str.length - 1; i >= 0; i--) { - const char = str[i] - if (char! in parenthesis) { - continue - } - - if (char in closeToOpen) { - stack.push(char) - if (cnt && cnt === resultSet.length) { - return resultSet - } - } else if (char in openToClose) { - if (stack.length !== 0 && stack[stack.length - 1] === openToClose[char]) { - stack.pop() - } else { - resultSet.push({ char: char, strOffset: i }) - } - } - } - - return resultSet -} - -function nonClosedClosingParen(str: string, cnt?: number): { char: string; strOffset: number }[] { - const resultSet: { char: string; strOffset: number }[] = [] - const stack: string[] = [] - - for (let i = 0; i < str.length; i++) { - const char = str[i] - if (char! in parenthesis) { - continue - } - - if (char in openToClose) { - stack.push(char) - if (cnt && cnt === resultSet.length) { - return resultSet - } - } else if (char in closeToOpen) { - if (stack.length !== 0 && stack[stack.length - 1] === closeToOpen[char]) { - stack.pop() - } else { - resultSet.push({ char: char, strOffset: i }) - } - } - } - - return resultSet -} diff --git a/packages/core/src/codewhisperer/util/commonUtil.ts b/packages/core/src/codewhisperer/util/commonUtil.ts index d2df78f1369..729d3b7ed12 100644 --- a/packages/core/src/codewhisperer/util/commonUtil.ts +++ b/packages/core/src/codewhisperer/util/commonUtil.ts @@ -3,80 +3,18 @@ * SPDX-License-Identifier: Apache-2.0 */ -import * as vscode from 'vscode' -import * as semver from 'semver' import { distance } from 'fastest-levenshtein' import { getInlineSuggestEnabled } from '../../shared/utilities/editorUtilities' -import { - AWSTemplateCaseInsensitiveKeyWords, - AWSTemplateKeyWords, - JsonConfigFileNamingConvention, -} from '../models/constants' export function getLocalDatetime() { const timezone = Intl.DateTimeFormat().resolvedOptions().timeZone return new Date().toLocaleString([], { timeZone: timezone }) } -export function asyncCallWithTimeout(asyncPromise: Promise, message: string, timeLimit: number): Promise { - let timeoutHandle: NodeJS.Timeout - const timeoutPromise = new Promise((_resolve, reject) => { - timeoutHandle = setTimeout(() => reject(new Error(message)), timeLimit) - }) - return Promise.race([asyncPromise, timeoutPromise]).then((result) => { - clearTimeout(timeoutHandle) - return result as T - }) -} - export function isInlineCompletionEnabled() { return getInlineSuggestEnabled() } -// This is the VS Code version that started to have regressions in inline completion API -export function isVscHavingRegressionInlineCompletionApi() { - return semver.gte(vscode.version, '1.78.0') && getInlineSuggestEnabled() -} - -export function getFileExt(languageId: string) { - switch (languageId) { - case 'java': - return '.java' - case 'python': - return '.py' - default: - break - } - return undefined -} - -/** - * Returns the longest overlap between the Suffix of firstString and Prefix of second string - * getPrefixSuffixOverlap("adwg31", "31ggrs") = "31" - */ -export function getPrefixSuffixOverlap(firstString: string, secondString: string) { - let i = Math.min(firstString.length, secondString.length) - while (i > 0) { - if (secondString.slice(0, i) === firstString.slice(-i)) { - break - } - i-- - } - return secondString.slice(0, i) -} - -export function checkLeftContextKeywordsForJson(fileName: string, leftFileContent: string, language: string): boolean { - if ( - language === 'json' && - !AWSTemplateKeyWords.some((substring) => leftFileContent.includes(substring)) && - !AWSTemplateCaseInsensitiveKeyWords.some((substring) => leftFileContent.toLowerCase().includes(substring)) && - !JsonConfigFileNamingConvention.has(fileName.toLowerCase()) - ) { - return true - } - return false -} - // With edit distance, complicate usermodification can be considered as simple edit(add, delete, replace), // and thus the unmodified part of recommendation length can be deducted/approximated // ex. (modified > original): originalRecom: foo -> modifiedRecom: fobarbarbaro, distance = 9, delta = 12 - 9 = 3 diff --git a/packages/core/src/codewhisperer/util/customizationUtil.ts b/packages/core/src/codewhisperer/util/customizationUtil.ts index 600317c53e0..04bb85d7a43 100644 --- a/packages/core/src/codewhisperer/util/customizationUtil.ts +++ b/packages/core/src/codewhisperer/util/customizationUtil.ts @@ -89,7 +89,14 @@ export const onProfileChangedListener: (event: ProfileChangedEvent) => any = asy */ export const getNewCustomizations = (availableCustomizations: Customization[]) => { const persistedCustomizations = getPersistedCustomizations() - return availableCustomizations.filter((c) => !persistedCustomizations.map((p) => p.arn).includes(c.arn)) + const newCustomizations = availableCustomizations.filter( + (c) => + !persistedCustomizations + .flat() + .map((p) => p.arn) + .includes(c.arn) + ) + return newCustomizations } export async function notifyNewCustomizations() { @@ -157,20 +164,11 @@ export const baseCustomization = { * @returns customization selected by users, `baseCustomization` if none is selected */ export const getSelectedCustomization = (): Customization => { - if ( - !AuthUtil.instance.isCustomizationFeatureEnabled || - !AuthUtil.instance.isValidEnterpriseSsoInUse() || - !AuthUtil.instance.conn - ) { + if (!AuthUtil.instance.isCustomizationFeatureEnabled || !AuthUtil.instance.isIdcConnection()) { return baseCustomization } - const selectedCustomizationArr = globals.globalState.tryGet<{ [label: string]: Customization }>( - 'CODEWHISPERER_SELECTED_CUSTOMIZATION', - Object, - {} - ) - const selectedCustomization = selectedCustomizationArr[AuthUtil.instance.conn.label] + const selectedCustomization = globals.globalState.getAmazonQCustomization(AuthUtil.instance.profileName) if (selectedCustomization && selectedCustomization.name !== '') { return selectedCustomization @@ -187,7 +185,7 @@ export const getSelectedCustomization = (): Customization => { * 2. the override customization arn is different from the previous override customization if any. The purpose is to only do override once on users' behalf. */ export const setSelectedCustomization = async (customization: Customization, isOverride: boolean = false) => { - if (!AuthUtil.instance.isValidEnterpriseSsoInUse() || !AuthUtil.instance.conn) { + if (!AuthUtil.instance.isIdcConnection()) { return } if (isOverride) { @@ -196,15 +194,10 @@ export const setSelectedCustomization = async (customization: Customization, isO return } } - const selectedCustomizationObj = globals.globalState.tryGet<{ [label: string]: Customization }>( - 'CODEWHISPERER_SELECTED_CUSTOMIZATION', - Object, - {} - ) - selectedCustomizationObj[AuthUtil.instance.conn.label] = customization - getLogger().debug(`Selected customization ${customization.name} for ${AuthUtil.instance.conn.label}`) - await globals.globalState.update('CODEWHISPERER_SELECTED_CUSTOMIZATION', selectedCustomizationObj) + await globals.globalState.update('CODEWHISPERER_SELECTED_CUSTOMIZATION', customization) + getLogger().debug(`Selected customization ${customization.name} for ${AuthUtil.instance.profileName}`) + if (isOverride) { await globals.globalState.update('aws.amazonq.customization.overrideV2', customization.arn) } @@ -216,28 +209,18 @@ export const setSelectedCustomization = async (customization: Customization, isO } export const getPersistedCustomizations = (): Customization[] => { - if (!AuthUtil.instance.isValidEnterpriseSsoInUse() || !AuthUtil.instance.conn) { + if (!AuthUtil.instance.isIdcConnection()) { return [] } - const persistedCustomizationsObj = globals.globalState.tryGet<{ [label: string]: Customization[] }>( - 'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', - Object, - {} - ) - return persistedCustomizationsObj[AuthUtil.instance.conn.label] || [] + return globals.globalState.getAmazonQCachedCustomization(AuthUtil.instance.profileName) } export const setPersistedCustomizations = async (customizations: Customization[]) => { - if (!AuthUtil.instance.isValidEnterpriseSsoInUse() || !AuthUtil.instance.conn) { + if (!AuthUtil.instance.isIdcConnection()) { return } - const persistedCustomizationsObj = globals.globalState.tryGet<{ [label: string]: Customization[] }>( - 'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', - Object, - {} - ) - persistedCustomizationsObj[AuthUtil.instance.conn.label] = customizations - await globals.globalState.update('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', persistedCustomizationsObj) + + await globals.globalState.update('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', customizations) } export const getNewCustomizationsAvailable = () => { diff --git a/packages/core/src/codewhisperer/util/editorContext.ts b/packages/core/src/codewhisperer/util/editorContext.ts deleted file mode 100644 index 95df5eb509a..00000000000 --- a/packages/core/src/codewhisperer/util/editorContext.ts +++ /dev/null @@ -1,425 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import * as codewhispererClient from '../client/codewhisperer' -import * as path from 'path' -import * as CodeWhispererConstants from '../models/constants' -import { getTabSizeSetting } from '../../shared/utilities/editorUtilities' -import { truncate } from '../../shared/utilities/textUtilities' -import { getLogger } from '../../shared/logger/logger' -import { runtimeLanguageContext } from './runtimeLanguageContext' -import { fetchSupplementalContext } from './supplementalContext/supplementalContextUtil' -import { editorStateMaxLength, supplementalContextTimeoutInMs } from '../models/constants' -import { getSelectedCustomization } from './customizationUtil' -import { selectFrom } from '../../shared/utilities/tsUtils' -import { checkLeftContextKeywordsForJson } from './commonUtil' -import { CodeWhispererSupplementalContext } from '../models/model' -import { getOptOutPreference } from '../../shared/telemetry/util' -import { indent } from '../../shared/utilities/textUtilities' -import { isInDirectory } from '../../shared/filesystemUtilities' -import { AuthUtil } from './authUtil' -import { predictionTracker } from '../nextEditPrediction/activation' - -let tabSize: number = getTabSizeSetting() - -function getEnclosingNotebook(editor: vscode.TextEditor): vscode.NotebookDocument | undefined { - // For notebook cells, find the existing notebook with a cell that matches the current editor. - return vscode.workspace.notebookDocuments.find( - (nb) => - nb.notebookType === 'jupyter-notebook' && nb.getCells().some((cell) => cell.document === editor.document) - ) -} - -export function getNotebookContext( - notebook: vscode.NotebookDocument, - editor: vscode.TextEditor, - languageName: string, - caretLeftFileContext: string, - caretRightFileContext: string -) { - // Expand the context for a cell inside of a noteboo with whatever text fits from the preceding and subsequent cells - const allCells = notebook.getCells() - const cellIndex = allCells.findIndex((cell) => cell.document === editor.document) - // Extract text from prior cells if there is enough room in left file context - if (caretLeftFileContext.length < CodeWhispererConstants.charactersLimit - 1) { - const leftCellsText = getNotebookCellsSliceContext( - allCells.slice(0, cellIndex), - CodeWhispererConstants.charactersLimit - (caretLeftFileContext.length + 1), - languageName, - true - ) - if (leftCellsText.length > 0) { - caretLeftFileContext = addNewlineIfMissing(leftCellsText) + caretLeftFileContext - } - } - // Extract text from subsequent cells if there is enough room in right file context - if (caretRightFileContext.length < CodeWhispererConstants.charactersLimit - 1) { - const rightCellsText = getNotebookCellsSliceContext( - allCells.slice(cellIndex + 1), - CodeWhispererConstants.charactersLimit - (caretRightFileContext.length + 1), - languageName, - false - ) - if (rightCellsText.length > 0) { - caretRightFileContext = addNewlineIfMissing(caretRightFileContext) + rightCellsText - } - } - return { caretLeftFileContext, caretRightFileContext } -} - -export function getNotebookCellContext(cell: vscode.NotebookCell, referenceLanguage?: string): string { - // Extract the text verbatim if the cell is code and the cell has the same language. - // Otherwise, add the correct comment string for the reference language - const cellText = cell.document.getText() - if ( - cell.kind === vscode.NotebookCellKind.Markup || - (runtimeLanguageContext.normalizeLanguage(cell.document.languageId) ?? cell.document.languageId) !== - referenceLanguage - ) { - const commentPrefix = runtimeLanguageContext.getSingleLineCommentPrefix(referenceLanguage) - if (commentPrefix === '') { - return cellText - } - return cell.document - .getText() - .split('\n') - .map((line) => `${commentPrefix}${line}`) - .join('\n') - } - return cellText -} - -export function getNotebookCellsSliceContext( - cells: vscode.NotebookCell[], - maxLength: number, - referenceLanguage: string, - fromStart: boolean -): string { - // Extract context from array of notebook cells that fits inside `maxLength` characters, - // from either the start or the end of the array. - let output: string[] = [] - if (!fromStart) { - cells = cells.reverse() - } - cells.some((cell) => { - const cellText = addNewlineIfMissing(getNotebookCellContext(cell, referenceLanguage)) - if (cellText.length > 0) { - if (cellText.length >= maxLength) { - if (fromStart) { - output.push(cellText.substring(0, maxLength)) - } else { - output.push(cellText.substring(cellText.length - maxLength)) - } - return true - } - output.push(cellText) - maxLength -= cellText.length - } - }) - if (!fromStart) { - output = output.reverse() - } - return output.join('') -} - -export function addNewlineIfMissing(text: string): string { - if (text.length > 0 && !text.endsWith('\n')) { - text += '\n' - } - return text -} - -export function extractContextForCodeWhisperer(editor: vscode.TextEditor): codewhispererClient.FileContext { - const document = editor.document - const curPos = editor.selection.active - const offset = document.offsetAt(curPos) - - let caretLeftFileContext = editor.document.getText( - new vscode.Range( - document.positionAt(offset - CodeWhispererConstants.charactersLimit), - document.positionAt(offset) - ) - ) - let caretRightFileContext = editor.document.getText( - new vscode.Range( - document.positionAt(offset), - document.positionAt(offset + CodeWhispererConstants.charactersLimit) - ) - ) - let languageName = 'plaintext' - if (!checkLeftContextKeywordsForJson(document.fileName, caretLeftFileContext, editor.document.languageId)) { - languageName = runtimeLanguageContext.resolveLang(editor.document) - } - if (editor.document.uri.scheme === 'vscode-notebook-cell') { - const notebook = getEnclosingNotebook(editor) - if (notebook) { - ;({ caretLeftFileContext, caretRightFileContext } = getNotebookContext( - notebook, - editor, - languageName, - caretLeftFileContext, - caretRightFileContext - )) - } - } - - return { - fileUri: editor.document.uri.toString().substring(0, CodeWhispererConstants.filenameCharsLimit), - filename: getFileRelativePath(editor), - programmingLanguage: { - languageName: languageName, - }, - leftFileContent: caretLeftFileContext, - rightFileContent: caretRightFileContext, - } as codewhispererClient.FileContext -} - -export function getFileName(editor: vscode.TextEditor): string { - const fileName = path.basename(editor.document.fileName) - return fileName.substring(0, CodeWhispererConstants.filenameCharsLimit) -} - -export function getFileRelativePath(editor: vscode.TextEditor): string { - const fileName = path.basename(editor.document.fileName) - let relativePath = '' - const workspaceFolder = vscode.workspace.getWorkspaceFolder(editor.document.uri) - if (!workspaceFolder) { - relativePath = fileName - } else { - const workspacePath = workspaceFolder.uri.fsPath - const filePath = editor.document.uri.fsPath - relativePath = path.relative(workspacePath, filePath) - } - // For notebook files, we want to use the programming language for each cell for the code suggestions, so change - // the filename sent in the request to reflect that language - if (relativePath.endsWith('.ipynb')) { - const fileExtension = runtimeLanguageContext.getLanguageExtensionForNotebook(editor.document.languageId) - if (fileExtension !== undefined) { - const filenameWithNewExtension = relativePath.substring(0, relativePath.length - 5) + fileExtension - return filenameWithNewExtension.substring(0, CodeWhispererConstants.filenameCharsLimit) - } - } - return relativePath.substring(0, CodeWhispererConstants.filenameCharsLimit) -} - -async function getWorkspaceId(editor: vscode.TextEditor): Promise { - try { - const workspaceIds: { workspaces: { workspaceRoot: string; workspaceId: string }[] } = - await vscode.commands.executeCommand('aws.amazonq.getWorkspaceId') - for (const item of workspaceIds.workspaces) { - const path = vscode.Uri.parse(item.workspaceRoot).fsPath - if (isInDirectory(path, editor.document.uri.fsPath)) { - return item.workspaceId - } - } - } catch (err) { - getLogger().warn(`No workspace id found ${err}`) - } - return undefined -} - -export async function buildListRecommendationRequest( - editor: vscode.TextEditor, - nextToken: string, - allowCodeWithReference: boolean -): Promise<{ - request: codewhispererClient.ListRecommendationsRequest - supplementalMetadata: CodeWhispererSupplementalContext | undefined -}> { - const fileContext = extractContextForCodeWhisperer(editor) - - const tokenSource = new vscode.CancellationTokenSource() - setTimeout(() => { - tokenSource.cancel() - }, supplementalContextTimeoutInMs) - - const supplementalContexts = await fetchSupplementalContext(editor, tokenSource.token) - - logSupplementalContext(supplementalContexts) - - // Get predictionSupplementalContext from PredictionTracker - let predictionSupplementalContext: codewhispererClient.SupplementalContext[] = [] - if (predictionTracker) { - predictionSupplementalContext = await predictionTracker.generatePredictionSupplementalContext() - } - - const selectedCustomization = getSelectedCustomization() - const completionSupplementalContext: codewhispererClient.SupplementalContext[] = supplementalContexts - ? supplementalContexts.supplementalContextItems.map((v) => { - return selectFrom(v, 'content', 'filePath') - }) - : [] - - const profile = AuthUtil.instance.regionProfileManager.activeRegionProfile - - const editorState = getEditorState(editor, fileContext) - - // Combine inline and prediction supplemental contexts - const finalSupplementalContext = completionSupplementalContext.concat(predictionSupplementalContext) - return { - request: { - fileContext: fileContext, - nextToken: nextToken, - referenceTrackerConfiguration: { - recommendationsWithReferences: allowCodeWithReference ? 'ALLOW' : 'BLOCK', - }, - supplementalContexts: finalSupplementalContext, - editorState: editorState, - maxResults: CodeWhispererConstants.maxRecommendations, - customizationArn: selectedCustomization.arn === '' ? undefined : selectedCustomization.arn, - optOutPreference: getOptOutPreference(), - workspaceId: await getWorkspaceId(editor), - profileArn: profile?.arn, - }, - supplementalMetadata: supplementalContexts, - } -} - -export async function buildGenerateRecommendationRequest(editor: vscode.TextEditor): Promise<{ - request: codewhispererClient.GenerateRecommendationsRequest - supplementalMetadata: CodeWhispererSupplementalContext | undefined -}> { - const fileContext = extractContextForCodeWhisperer(editor) - - const tokenSource = new vscode.CancellationTokenSource() - // the supplement context fetch mechanisms each has a timeout of supplementalContextTimeoutInMs - // adding 10 ms for overall timeout as buffer - setTimeout(() => { - tokenSource.cancel() - }, supplementalContextTimeoutInMs + 10) - const supplementalContexts = await fetchSupplementalContext(editor, tokenSource.token) - - logSupplementalContext(supplementalContexts) - - return { - request: { - fileContext: fileContext, - maxResults: CodeWhispererConstants.maxRecommendations, - supplementalContexts: supplementalContexts?.supplementalContextItems ?? [], - }, - supplementalMetadata: supplementalContexts, - } -} - -export function validateRequest( - req: codewhispererClient.ListRecommendationsRequest | codewhispererClient.GenerateRecommendationsRequest -): boolean { - const isLanguageNameValid = - req.fileContext.programmingLanguage.languageName !== undefined && - req.fileContext.programmingLanguage.languageName.length >= 1 && - req.fileContext.programmingLanguage.languageName.length <= 128 && - (runtimeLanguageContext.isLanguageSupported(req.fileContext.programmingLanguage.languageName) || - runtimeLanguageContext.isFileFormatSupported( - req.fileContext.filename.substring(req.fileContext.filename.lastIndexOf('.') + 1) - )) - const isFileNameValid = !(req.fileContext.filename === undefined || req.fileContext.filename.length < 1) - const isFileContextValid = !( - req.fileContext.leftFileContent.length > CodeWhispererConstants.charactersLimit || - req.fileContext.rightFileContent.length > CodeWhispererConstants.charactersLimit - ) - if (isFileNameValid && isLanguageNameValid && isFileContextValid) { - return true - } - return false -} - -export function updateTabSize(val: number): void { - tabSize = val -} - -export function getTabSize(): number { - return tabSize -} - -export function getEditorState(editor: vscode.TextEditor, fileContext: codewhispererClient.FileContext): any { - try { - const cursorPosition = editor.selection.active - const cursorOffset = editor.document.offsetAt(cursorPosition) - const documentText = editor.document.getText() - - // Truncate if document content is too large (defined in constants.ts) - let fileText = documentText - if (documentText.length > editorStateMaxLength) { - const halfLength = Math.floor(editorStateMaxLength / 2) - - // Use truncate function to get the text around the cursor position - const leftPart = truncate(documentText.substring(0, cursorOffset), -halfLength, '') - const rightPart = truncate(documentText.substring(cursorOffset), halfLength, '') - - fileText = leftPart + rightPart - } - - return { - document: { - programmingLanguage: { - languageName: fileContext.programmingLanguage.languageName, - }, - relativeFilePath: fileContext.filename, - text: fileText, - }, - cursorState: { - position: { - line: editor.selection.active.line, - character: editor.selection.active.character, - }, - }, - } - } catch (error) { - getLogger().error(`Error generating editor state: ${error}`) - return undefined - } -} - -export function getLeftContext(editor: vscode.TextEditor, line: number): string { - let lineText = '' - try { - if (editor && editor.document.lineAt(line)) { - lineText = editor.document.lineAt(line).text - if (lineText.length > CodeWhispererConstants.contextPreviewLen) { - lineText = - '...' + - lineText.substring( - lineText.length - CodeWhispererConstants.contextPreviewLen - 1, - lineText.length - 1 - ) - } - } - } catch (error) { - getLogger().error(`Error when getting left context ${error}`) - } - - return lineText -} - -function logSupplementalContext(supplementalContext: CodeWhispererSupplementalContext | undefined) { - if (!supplementalContext) { - return - } - - let logString = indent( - `CodeWhispererSupplementalContext: - isUtg: ${supplementalContext.isUtg}, - isProcessTimeout: ${supplementalContext.isProcessTimeout}, - contentsLength: ${supplementalContext.contentsLength}, - latency: ${supplementalContext.latency} - strategy: ${supplementalContext.strategy}`, - 4, - true - ).trimStart() - - for (const [index, context] of supplementalContext.supplementalContextItems.entries()) { - logString += indent(`\nChunk ${index}:\n`, 4, true) - logString += indent( - `Path: ${context.filePath} - Length: ${context.content.length} - Score: ${context.score}`, - 8, - true - ) - } - - getLogger().debug(logString) -} diff --git a/packages/core/src/codewhisperer/util/getStartUrl.ts b/packages/core/src/codewhisperer/util/getStartUrl.ts index 0d9ca7617ff..f1db38f5f1f 100644 --- a/packages/core/src/codewhisperer/util/getStartUrl.ts +++ b/packages/core/src/codewhisperer/util/getStartUrl.ts @@ -28,9 +28,8 @@ export const getStartUrl = async () => { } export async function connectToEnterpriseSso(startUrl: string, region: Region['id']) { - let conn try { - conn = await AuthUtil.instance.connectToEnterpriseSso(startUrl, region) + await AuthUtil.instance.login(startUrl, region) } catch (e) { throw ToolkitError.chain(e, CodeWhispererConstants.failedToConnectIamIdentityCenter, { code: 'FailedToConnect', @@ -38,6 +37,4 @@ export async function connectToEnterpriseSso(startUrl: string, region: Region['i } vsCodeState.isFreeTierLimitReached = false await Commands.tryExecute('aws.amazonq.enableCodeSuggestions') - - return conn } diff --git a/packages/core/src/codewhisperer/util/globalStateUtil.ts b/packages/core/src/codewhisperer/util/globalStateUtil.ts deleted file mode 100644 index 55376a83546..00000000000 --- a/packages/core/src/codewhisperer/util/globalStateUtil.ts +++ /dev/null @@ -1,23 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import { vsCodeState } from '../models/model' - -export function resetIntelliSenseState( - isManualTriggerEnabled: boolean, - isAutomatedTriggerEnabled: boolean, - hasResponse: boolean -) { - /** - * Skip when CodeWhisperer service is turned off - */ - if (!isManualTriggerEnabled && !isAutomatedTriggerEnabled) { - return - } - - if (vsCodeState.isIntelliSenseActive && hasResponse) { - vsCodeState.isIntelliSenseActive = false - } -} diff --git a/packages/core/src/codewhisperer/util/showSsoPrompt.ts b/packages/core/src/codewhisperer/util/showSsoPrompt.ts index 13af4cf771f..b3d78654745 100644 --- a/packages/core/src/codewhisperer/util/showSsoPrompt.ts +++ b/packages/core/src/codewhisperer/util/showSsoPrompt.ts @@ -17,6 +17,7 @@ import { telemetry } from '../../shared/telemetry/telemetry' import { createBuilderIdItem, createSsoItem, createIamItem } from '../../auth/utils' import { Commands } from '../../shared/vscode/commands2' import { vsCodeState } from '../models/model' +import { builderIdRegion, builderIdStartUrl } from '../../auth/sso/constants' export const showCodeWhispererConnectionPrompt = async () => { const items = [createBuilderIdItem(), createSsoItem(), createCodeWhispererIamItem()] @@ -45,16 +46,13 @@ export const showCodeWhispererConnectionPrompt = async () => { export async function awsIdSignIn() { getLogger().info('selected AWS ID sign in') - let conn try { - conn = await AuthUtil.instance.connectToAwsBuilderId() + await AuthUtil.instance.login(builderIdStartUrl, builderIdRegion) } catch (e) { throw ToolkitError.chain(e, failedToConnectAwsBuilderId, { code: 'FailedToConnect' }) } vsCodeState.isFreeTierLimitReached = false await Commands.tryExecute('aws.amazonq.enableCodeSuggestions') - - return conn } export const createCodeWhispererIamItem = () => { diff --git a/packages/core/src/codewhisperer/util/supplementalContext/codeParsingUtil.ts b/packages/core/src/codewhisperer/util/supplementalContext/codeParsingUtil.ts deleted file mode 100644 index c73a2eebaa4..00000000000 --- a/packages/core/src/codewhisperer/util/supplementalContext/codeParsingUtil.ts +++ /dev/null @@ -1,130 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import path = require('path') -import { normalize } from '../../../shared/utilities/pathUtils' - -// TODO: functionExtractionPattern, classExtractionPattern, imposrtStatementRegex are not scalable and we will deprecate and remove the usage in the near future -export interface utgLanguageConfig { - extension: string - testFilenamePattern: RegExp[] - functionExtractionPattern?: RegExp - classExtractionPattern?: RegExp - importStatementRegExp?: RegExp -} - -export const utgLanguageConfigs: Record = { - // Java regexes are not working efficiently for class or function extraction - java: { - extension: '.java', - testFilenamePattern: [/^(.+)Test(\.java)$/, /(.+)Tests(\.java)$/, /Test(.+)(\.java)$/], - functionExtractionPattern: - /(?:(?:public|private|protected)\s+)(?:static\s+)?(?:[\w<>]+\s+)?(\w+)\s*\([^)]*\)\s*(?:(?:throws\s+\w+)?\s*)[{;]/gm, // TODO: Doesn't work for generice T functions. - classExtractionPattern: /(?<=^|\n)\s*public\s+class\s+(\w+)/gm, // TODO: Verify these. - importStatementRegExp: /import .*\.([a-zA-Z0-9]+);/, - }, - python: { - extension: '.py', - testFilenamePattern: [/^test_(.+)(\.py)$/, /^(.+)_test(\.py)$/], - functionExtractionPattern: /def\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\(/g, // Worked fine - classExtractionPattern: /^class\s+(\w+)\s*:/gm, - importStatementRegExp: /from (.*) import.*/, - }, - typescript: { - extension: '.ts', - testFilenamePattern: [/^(.+)\.test(\.ts|\.tsx)$/, /^(.+)\.spec(\.ts|\.tsx)$/], - }, - javascript: { - extension: '.js', - testFilenamePattern: [/^(.+)\.test(\.js|\.jsx)$/, /^(.+)\.spec(\.js|\.jsx)$/], - }, - typescriptreact: { - extension: '.tsx', - testFilenamePattern: [/^(.+)\.test(\.ts|\.tsx)$/, /^(.+)\.spec(\.ts|\.tsx)$/], - }, - javascriptreact: { - extension: '.jsx', - testFilenamePattern: [/^(.+)\.test(\.js|\.jsx)$/, /^(.+)\.spec(\.js|\.jsx)$/], - }, -} - -export function extractFunctions(fileContent: string, regex?: RegExp) { - if (!regex) { - return [] - } - const functionNames: string[] = [] - let match: RegExpExecArray | null - - while ((match = regex.exec(fileContent)) !== null) { - functionNames.push(match[1]) - } - return functionNames -} - -export function extractClasses(fileContent: string, regex?: RegExp) { - if (!regex) { - return [] - } - const classNames: string[] = [] - let match: RegExpExecArray | null - - while ((match = regex.exec(fileContent)) !== null) { - classNames.push(match[1]) - } - return classNames -} - -export function countSubstringMatches(arr1: string[], arr2: string[]): number { - let count = 0 - for (const str1 of arr1) { - for (const str2 of arr2) { - if (str2.toLowerCase().includes(str1.toLowerCase())) { - count++ - } - } - } - return count -} - -export async function isTestFile( - filePath: string, - languageConfig: { - languageId: vscode.TextDocument['languageId'] - fileContent?: string - } -): Promise { - const normalizedFilePath = normalize(filePath) - const pathContainsTest = - normalizedFilePath.includes('tests/') || - normalizedFilePath.includes('test/') || - normalizedFilePath.includes('tst/') - const fileNameMatchTestPatterns = isTestFileByName(normalizedFilePath, languageConfig.languageId) - - if (pathContainsTest || fileNameMatchTestPatterns) { - return true - } - - return false -} - -function isTestFileByName(filePath: string, language: vscode.TextDocument['languageId']): boolean { - const languageConfig = utgLanguageConfigs[language] - if (!languageConfig) { - // We have enabled the support only for python and Java for this check - // as we depend on Regex for this validation. - return false - } - const testFilenamePattern = languageConfig.testFilenamePattern - - const filename = path.basename(filePath) - for (const pattern of testFilenamePattern) { - if (pattern.test(filename)) { - return true - } - } - - return false -} diff --git a/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts b/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts deleted file mode 100644 index db1d7f312b2..00000000000 --- a/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts +++ /dev/null @@ -1,395 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as vscode from 'vscode' -import path = require('path') -import { BM25Document, BM25Okapi } from './rankBm25' -import { - crossFileContextConfig, - supplementalContextTimeoutInMs, - supplementalContextMaxTotalLength, -} from '../../models/constants' -import { isTestFile } from './codeParsingUtil' -import { getFileDistance } from '../../../shared/filesystemUtilities' -import { getOpenFilesInWindow } from '../../../shared/utilities/editorUtilities' -import { getLogger } from '../../../shared/logger/logger' -import { - CodeWhispererSupplementalContext, - CodeWhispererSupplementalContextItem, - SupplementalContextStrategy, -} from '../../models/model' -import { LspController } from '../../../amazonq/lsp/lspController' -import { waitUntil } from '../../../shared/utilities/timeoutUtils' -import { FeatureConfigProvider } from '../../../shared/featureConfig' -import fs from '../../../shared/fs/fs' - -type CrossFileSupportedLanguage = - | 'java' - | 'python' - | 'javascript' - | 'typescript' - | 'javascriptreact' - | 'typescriptreact' - -// TODO: ugly, can we make it prettier? like we have to manually type 'java', 'javascriptreact' which is error prone -// TODO: Move to another config file or constants file -// Supported language to its corresponding file ext -const supportedLanguageToDialects: Readonly>> = { - java: new Set(['.java']), - python: new Set(['.py']), - javascript: new Set(['.js', '.jsx']), - javascriptreact: new Set(['.js', '.jsx']), - typescript: new Set(['.ts', '.tsx']), - typescriptreact: new Set(['.ts', '.tsx']), -} - -function isCrossFileSupported(languageId: string): languageId is CrossFileSupportedLanguage { - return Object.keys(supportedLanguageToDialects).includes(languageId) -} - -interface Chunk { - fileName: string - content: string - nextContent: string - score?: number -} - -/** - * `none`: supplementalContext is not supported - * `opentabs`: opentabs_BM25 - * `codemap`: repomap + opentabs BM25 - * `bm25`: global_BM25 - * `default`: repomap + global_BM25 - */ -type SupplementalContextConfig = 'none' | 'opentabs' | 'codemap' | 'bm25' | 'default' - -export async function fetchSupplementalContextForSrc( - editor: vscode.TextEditor, - cancellationToken: vscode.CancellationToken -): Promise | undefined> { - const supplementalContextConfig = getSupplementalContextConfig(editor.document.languageId) - - // not supported case - if (supplementalContextConfig === 'none') { - return undefined - } - - // fallback to opentabs if projectContext timeout - const opentabsContextPromise = waitUntil( - async function () { - return await fetchOpentabsContext(editor, cancellationToken) - }, - { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false } - ) - - // opentabs context will use bm25 and users' open tabs to fetch supplemental context - if (supplementalContextConfig === 'opentabs') { - const supContext = (await opentabsContextPromise) ?? [] - return { - supplementalContextItems: supContext, - strategy: supContext.length === 0 ? 'empty' : 'opentabs', - } - } - - // codemap will use opentabs context plus repomap if it's present - if (supplementalContextConfig === 'codemap') { - let strategy: SupplementalContextStrategy = 'empty' - let hasCodemap: boolean = false - let hasOpentabs: boolean = false - const opentabsContextAndCodemap = await waitUntil( - async function () { - const result: CodeWhispererSupplementalContextItem[] = [] - const opentabsContext = await fetchOpentabsContext(editor, cancellationToken) - const codemap = await fetchProjectContext(editor, 'codemap') - - function addToResult(items: CodeWhispererSupplementalContextItem[]) { - for (const item of items) { - const curLen = result.reduce((acc, i) => acc + i.content.length, 0) - if (curLen + item.content.length < supplementalContextMaxTotalLength) { - result.push(item) - } - } - } - - if (codemap && codemap.length > 0) { - addToResult(codemap) - hasCodemap = true - } - - if (opentabsContext && opentabsContext.length > 0) { - addToResult(opentabsContext) - hasOpentabs = true - } - - return result - }, - { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false } - ) - - if (hasCodemap) { - strategy = 'codemap' - } else if (hasOpentabs) { - strategy = 'opentabs' - } else { - strategy = 'empty' - } - - return { - supplementalContextItems: opentabsContextAndCodemap ?? [], - strategy: strategy, - } - } - - // global bm25 without repomap - if (supplementalContextConfig === 'bm25') { - const projectBM25Promise = waitUntil( - async function () { - return await fetchProjectContext(editor, 'bm25') - }, - { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false } - ) - - const [projectContext, opentabsContext] = await Promise.all([projectBM25Promise, opentabsContextPromise]) - if (projectContext && projectContext.length > 0) { - return { - supplementalContextItems: projectContext, - strategy: 'bm25', - } - } - - const supContext = opentabsContext ?? [] - return { - supplementalContextItems: supContext, - strategy: supContext.length === 0 ? 'empty' : 'opentabs', - } - } - - // global bm25 with repomap - const projectContextAndCodemapPromise = waitUntil( - async function () { - return await fetchProjectContext(editor, 'default') - }, - { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false } - ) - - const [projectContext, opentabsContext] = await Promise.all([ - projectContextAndCodemapPromise, - opentabsContextPromise, - ]) - if (projectContext && projectContext.length > 0) { - return { - supplementalContextItems: projectContext, - strategy: 'default', - } - } - - return { - supplementalContextItems: opentabsContext ?? [], - strategy: 'opentabs', - } -} - -export async function fetchProjectContext( - editor: vscode.TextEditor, - target: 'default' | 'codemap' | 'bm25' -): Promise { - const inputChunkContent = getInputChunk(editor) - - const inlineProjectContext: { content: string; score: number; filePath: string }[] = - await LspController.instance.queryInlineProjectContext( - inputChunkContent.content, - editor.document.uri.fsPath, - target - ) - - return inlineProjectContext -} - -export async function fetchOpentabsContext( - editor: vscode.TextEditor, - cancellationToken: vscode.CancellationToken -): Promise { - const codeChunksCalculated = crossFileContextConfig.numberOfChunkToFetch - - // Step 1: Get relevant cross files to refer - const relevantCrossFilePaths = await getCrossFileCandidates(editor) - - // Step 2: Split files to chunks with upper bound on chunkCount - // We restrict the total number of chunks to improve on latency. - // Chunk linking is required as we want to pass the next chunk value for matched chunk. - let chunkList: Chunk[] = [] - for (const relevantFile of relevantCrossFilePaths) { - const chunks: Chunk[] = await splitFileToChunks(relevantFile, crossFileContextConfig.numberOfLinesEachChunk) - const linkedChunks = linkChunks(chunks) - chunkList.push(...linkedChunks) - if (chunkList.length >= codeChunksCalculated) { - break - } - } - - // it's required since chunkList.push(...) is likely giving us a list of size > 60 - chunkList = chunkList.slice(0, codeChunksCalculated) - - // Step 3: Generate Input chunk (10 lines left of cursor position) - // and Find Best K chunks w.r.t input chunk using BM25 - const inputChunk: Chunk = getInputChunk(editor) - const bestChunks: Chunk[] = findBestKChunkMatches(inputChunk, chunkList, crossFileContextConfig.topK) - - // Step 4: Transform best chunks to supplemental contexts - const supplementalContexts: CodeWhispererSupplementalContextItem[] = [] - let totalLength = 0 - for (const chunk of bestChunks) { - totalLength += chunk.nextContent.length - - if (totalLength > crossFileContextConfig.maximumTotalLength) { - break - } - - supplementalContexts.push({ - filePath: chunk.fileName, - content: chunk.nextContent, - score: chunk.score, - }) - } - - // DO NOT send code chunk with empty content - getLogger().debug(`CodeWhisperer finished fetching crossfile context out of ${relevantCrossFilePaths.length} files`) - return supplementalContexts -} - -function findBestKChunkMatches(chunkInput: Chunk, chunkReferences: Chunk[], k: number): Chunk[] { - const chunkContentList = chunkReferences.map((chunk) => chunk.content) - - // performBM25Scoring returns the output in a sorted order (descending of scores) - const top3: BM25Document[] = new BM25Okapi(chunkContentList).topN(chunkInput.content, crossFileContextConfig.topK) - - return top3.map((doc) => { - // reference to the original metadata since BM25.top3 will sort the result - const chunkIndex = doc.index - const chunkReference = chunkReferences[chunkIndex] - return { - content: chunkReference.content, - fileName: chunkReference.fileName, - nextContent: chunkReference.nextContent, - score: doc.score, - } - }) -} - -/* This extract 10 lines to the left of the cursor from trigger file. - * This will be the inputquery to bm25 matching against list of cross-file chunks - */ -function getInputChunk(editor: vscode.TextEditor) { - const chunkSize = crossFileContextConfig.numberOfLinesEachChunk - const cursorPosition = editor.selection.active - const startLine = Math.max(cursorPosition.line - chunkSize, 0) - const endLine = Math.max(cursorPosition.line - 1, 0) - const inputChunkContent = editor.document.getText( - new vscode.Range(startLine, 0, endLine, editor.document.lineAt(endLine).text.length) - ) - const inputChunk: Chunk = { fileName: editor.document.fileName, content: inputChunkContent, nextContent: '' } - return inputChunk -} - -/** - * Util to decide if we need to fetch crossfile context since CodeWhisperer CrossFile Context feature is gated by userGroup and language level - * @param languageId: VSCode language Identifier - * @returns specifically returning undefined if the langueage is not supported, - * otherwise true/false depending on if the language is fully supported or not belonging to the user group - */ -function getSupplementalContextConfig(languageId: vscode.TextDocument['languageId']): SupplementalContextConfig { - if (!isCrossFileSupported(languageId)) { - return 'none' - } - - const group = FeatureConfigProvider.instance.getProjectContextGroup() - switch (group) { - default: - return 'codemap' - } -} - -/** - * This linking is required from science experimentations to pass the next contnet chunk - * when a given chunk context passes the match in BM25. - * Special handling is needed for last(its next points to its own) and first chunk - */ -export function linkChunks(chunks: Chunk[]) { - const updatedChunks: Chunk[] = [] - - // This additional chunk is needed to create a next pointer to chunk 0. - const firstChunk = chunks[0] - const firstChunkSubContent = firstChunk.content.split('\n').slice(0, 3).join('\n').trimEnd() - const newFirstChunk = { - fileName: firstChunk.fileName, - content: firstChunkSubContent, - nextContent: firstChunk.content, - } - updatedChunks.push(newFirstChunk) - - const n = chunks.length - for (let i = 0; i < n; i++) { - const chunk = chunks[i] - const nextChunk = i < n - 1 ? chunks[i + 1] : chunk - - chunk.nextContent = nextChunk.content - updatedChunks.push(chunk) - } - - return updatedChunks -} - -export async function splitFileToChunks(filePath: string, chunkSize: number): Promise { - const chunks: Chunk[] = [] - - const fileContent = (await fs.readFileText(filePath)).trimEnd() - const lines = fileContent.split('\n') - - for (let i = 0; i < lines.length; i += chunkSize) { - const chunkContent = lines.slice(i, Math.min(i + chunkSize, lines.length)).join('\n') - const chunk = { fileName: filePath, content: chunkContent.trimEnd(), nextContent: '' } - chunks.push(chunk) - } - return chunks -} - -/** - * This function will return relevant cross files sorted by file distance for the given editor file - * by referencing open files, imported files and same package files. - */ -export async function getCrossFileCandidates(editor: vscode.TextEditor): Promise { - const targetFile = editor.document.uri.fsPath - const language = editor.document.languageId as CrossFileSupportedLanguage - const dialects = supportedLanguageToDialects[language] - - /** - * Consider a file which - * 1. is different from the target - * 2. has the same file extension or it's one of the dialect of target file (e.g .js vs. .jsx) - * 3. is not a test file - */ - const unsortedCandidates = await getOpenFilesInWindow(async (candidateFile) => { - return ( - targetFile !== candidateFile && - (path.extname(targetFile) === path.extname(candidateFile) || - (dialects && dialects.has(path.extname(candidateFile)))) && - !(await isTestFile(candidateFile, { languageId: language })) - ) - }) - - return unsortedCandidates - .map((candidate) => { - return { - file: candidate, - fileDistance: getFileDistance(targetFile, candidate), - } - }) - .sort((file1, file2) => { - return file1.fileDistance - file2.fileDistance - }) - .map((fileToDistance) => { - return fileToDistance.file - }) -} diff --git a/packages/core/src/codewhisperer/util/supplementalContext/rankBm25.ts b/packages/core/src/codewhisperer/util/supplementalContext/rankBm25.ts deleted file mode 100644 index a2c77e0b10f..00000000000 --- a/packages/core/src/codewhisperer/util/supplementalContext/rankBm25.ts +++ /dev/null @@ -1,137 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -// Implementation inspired by https://github.com/dorianbrown/rank_bm25/blob/990470ebbe6b28c18216fd1a8b18fe7446237dd6/rank_bm25.py#L52 - -export interface BM25Document { - content: string - /** The score that the document receives. */ - score: number - - index: number -} - -export abstract class BM25 { - protected readonly corpusSize: number - protected readonly avgdl: number - protected readonly idf: Map = new Map() - protected readonly docLen: number[] = [] - protected readonly docFreqs: Map[] = [] - protected readonly nd: Map = new Map() - - constructor( - protected readonly corpus: string[], - protected readonly tokenizer: (str: string) => string[] = defaultTokenizer, - protected readonly k1: number, - protected readonly b: number, - protected readonly epsilon: number - ) { - this.corpusSize = corpus.length - - let numDoc = 0 - for (const document of corpus.map((document) => { - return tokenizer(document) - })) { - this.docLen.push(document.length) - numDoc += document.length - - const frequencies = new Map() - for (const word of document) { - frequencies.set(word, (frequencies.get(word) || 0) + 1) - } - this.docFreqs.push(frequencies) - - for (const [word, _] of frequencies.entries()) { - this.nd.set(word, (this.nd.get(word) || 0) + 1) - } - } - - this.avgdl = numDoc / this.corpusSize - - this.calIdf(this.nd) - } - - abstract calIdf(nd: Map): void - - abstract score(query: string): BM25Document[] - - topN(query: string, n: number): BM25Document[] { - const notSorted = this.score(query) - const sorted = notSorted.sort((a, b) => b.score - a.score) - return sorted.slice(0, Math.min(n, sorted.length)) - } -} - -export class BM25Okapi extends BM25 { - constructor(corpus: string[], tokenizer: (str: string) => string[] = defaultTokenizer) { - super(corpus, tokenizer, 1.5, 0.75, 0.25) - } - - calIdf(nd: Map): void { - let idfSum = 0 - - const negativeIdfs: string[] = [] - for (const [word, freq] of nd) { - const idf = Math.log(this.corpusSize - freq + 0.5) - Math.log(freq + 0.5) - this.idf.set(word, idf) - idfSum += idf - - if (idf < 0) { - negativeIdfs.push(word) - } - } - - const averageIdf = idfSum / this.idf.size - const eps = this.epsilon * averageIdf - for (const word of negativeIdfs) { - this.idf.set(word, eps) - } - } - - score(query: string): BM25Document[] { - const queryWords = defaultTokenizer(query) - return this.docFreqs.map((docFreq, index) => { - let score = 0 - for (const [_, queryWord] of queryWords.entries()) { - const queryWordFreqForDocument = docFreq.get(queryWord) || 0 - const numerator = (this.idf.get(queryWord) || 0.0) * queryWordFreqForDocument * (this.k1 + 1) - const denominator = - queryWordFreqForDocument + this.k1 * (1 - this.b + (this.b * this.docLen[index]) / this.avgdl) - - score += numerator / denominator - } - - return { - content: this.corpus[index], - score: score, - index: index, - } - }) - } -} - -// TODO: This is a very simple tokenizer, we want to replace this by more sophisticated one. -function defaultTokenizer(content: string): string[] { - const regex = /\w+/g - const words = content.split(' ') - const result = [] - for (const word of words) { - const wordList = findAll(word, regex) - result.push(...wordList) - } - - return result -} - -function findAll(str: string, re: RegExp): string[] { - let match: RegExpExecArray | null - const matches: string[] = [] - - while ((match = re.exec(str)) !== null) { - matches.push(match[0]) - } - - return matches -} diff --git a/packages/core/src/codewhisperer/util/supplementalContext/supplementalContextUtil.ts b/packages/core/src/codewhisperer/util/supplementalContext/supplementalContextUtil.ts deleted file mode 100644 index bd214ace44e..00000000000 --- a/packages/core/src/codewhisperer/util/supplementalContext/supplementalContextUtil.ts +++ /dev/null @@ -1,137 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import { fetchSupplementalContextForTest } from './utgUtils' -import { fetchSupplementalContextForSrc } from './crossFileContextUtil' -import { isTestFile } from './codeParsingUtil' -import * as vscode from 'vscode' -import { CancellationError } from '../../../shared/utilities/timeoutUtils' -import { ToolkitError } from '../../../shared/errors' -import { getLogger } from '../../../shared/logger/logger' -import { CodeWhispererSupplementalContext } from '../../models/model' -import * as os from 'os' -import { crossFileContextConfig } from '../../models/constants' - -export async function fetchSupplementalContext( - editor: vscode.TextEditor, - cancellationToken: vscode.CancellationToken -): Promise { - const timesBeforeFetching = performance.now() - - const isUtg = await isTestFile(editor.document.uri.fsPath, { - languageId: editor.document.languageId, - fileContent: editor.document.getText(), - }) - - let supplementalContextPromise: Promise< - Pick | undefined - > - - if (isUtg) { - supplementalContextPromise = fetchSupplementalContextForTest(editor, cancellationToken) - } else { - supplementalContextPromise = fetchSupplementalContextForSrc(editor, cancellationToken) - } - - return supplementalContextPromise - .then((value) => { - if (value) { - const resBeforeTruncation = { - isUtg: isUtg, - isProcessTimeout: false, - supplementalContextItems: value.supplementalContextItems.filter( - (item) => item.content.trim().length !== 0 - ), - contentsLength: value.supplementalContextItems.reduce((acc, curr) => acc + curr.content.length, 0), - latency: performance.now() - timesBeforeFetching, - strategy: value.strategy, - } - - return truncateSuppelementalContext(resBeforeTruncation) - } else { - return undefined - } - }) - .catch((err) => { - if (err instanceof ToolkitError && err.cause instanceof CancellationError) { - return { - isUtg: isUtg, - isProcessTimeout: true, - supplementalContextItems: [], - contentsLength: 0, - latency: performance.now() - timesBeforeFetching, - strategy: 'empty', - } - } else { - getLogger().error( - `Fail to fetch supplemental context for target file ${editor.document.fileName}: ${err}` - ) - return undefined - } - }) -} - -/** - * Requirement - * - Maximum 5 supplemental context. - * - Each chunk can't exceed 10240 characters - * - Sum of all chunks can't exceed 20480 characters - */ -export function truncateSuppelementalContext( - context: CodeWhispererSupplementalContext -): CodeWhispererSupplementalContext { - let c = context.supplementalContextItems.map((item) => { - if (item.content.length > crossFileContextConfig.maxLengthEachChunk) { - return { - ...item, - content: truncateLineByLine(item.content, crossFileContextConfig.maxLengthEachChunk), - } - } else { - return item - } - }) - - if (c.length > crossFileContextConfig.maxContextCount) { - c = c.slice(0, crossFileContextConfig.maxContextCount) - } - - let curTotalLength = c.reduce((acc, cur) => { - return acc + cur.content.length - }, 0) - while (curTotalLength >= 20480 && c.length - 1 >= 0) { - const last = c[c.length - 1] - c = c.slice(0, -1) - curTotalLength -= last.content.length - } - - return { - ...context, - supplementalContextItems: c, - contentsLength: curTotalLength, - } -} - -export function truncateLineByLine(input: string, l: number): string { - const maxLength = l > 0 ? l : -1 * l - if (input.length === 0) { - return '' - } - - const shouldAddNewLineBack = input.endsWith(os.EOL) - let lines = input.trim().split(os.EOL) - let curLen = input.length - while (curLen > maxLength && lines.length - 1 >= 0) { - const last = lines[lines.length - 1] - lines = lines.slice(0, -1) - curLen -= last.length + 1 - } - - const r = lines.join(os.EOL) - if (shouldAddNewLineBack) { - return r + os.EOL - } else { - return r - } -} diff --git a/packages/core/src/codewhisperer/util/supplementalContext/utgUtils.ts b/packages/core/src/codewhisperer/util/supplementalContext/utgUtils.ts deleted file mode 100644 index 0d33969773e..00000000000 --- a/packages/core/src/codewhisperer/util/supplementalContext/utgUtils.ts +++ /dev/null @@ -1,229 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import * as path from 'path' -import { fs } from '../../../shared/fs/fs' -import * as vscode from 'vscode' -import { - countSubstringMatches, - extractClasses, - extractFunctions, - isTestFile, - utgLanguageConfig, - utgLanguageConfigs, -} from './codeParsingUtil' -import { ToolkitError } from '../../../shared/errors' -import { supplemetalContextFetchingTimeoutMsg } from '../../models/constants' -import { CancellationError } from '../../../shared/utilities/timeoutUtils' -import { utgConfig } from '../../models/constants' -import { getOpenFilesInWindow } from '../../../shared/utilities/editorUtilities' -import { getLogger } from '../../../shared/logger/logger' -import { CodeWhispererSupplementalContext, CodeWhispererSupplementalContextItem, UtgStrategy } from '../../models/model' - -const utgSupportedLanguages: vscode.TextDocument['languageId'][] = ['java', 'python'] - -type UtgSupportedLanguage = (typeof utgSupportedLanguages)[number] - -function isUtgSupportedLanguage(languageId: vscode.TextDocument['languageId']): languageId is UtgSupportedLanguage { - return utgSupportedLanguages.includes(languageId) -} - -export function shouldFetchUtgContext(languageId: vscode.TextDocument['languageId']): boolean | undefined { - if (!isUtgSupportedLanguage(languageId)) { - return undefined - } - - return languageId === 'java' -} - -/** - * This function attempts to find a focal file for the given trigger file. - * Attempt 1: If naming patterns followed correctly, source file can be found by name referencing. - * Attempt 2: Compare the function and class names of trigger file and all other open files in editor - * to find the closest match. - * Once found the focal file, we split it into multiple pieces as supplementalContext. - * @param editor - * @returns - */ -export async function fetchSupplementalContextForTest( - editor: vscode.TextEditor, - cancellationToken: vscode.CancellationToken -): Promise | undefined> { - const shouldProceed = shouldFetchUtgContext(editor.document.languageId) - - if (!shouldProceed) { - return shouldProceed === undefined ? undefined : { supplementalContextItems: [], strategy: 'empty' } - } - - const languageConfig = utgLanguageConfigs[editor.document.languageId] - - // TODO (Metrics): 1. Total number of calls to fetchSupplementalContextForTest - throwIfCancelled(cancellationToken) - - let crossSourceFile = await findSourceFileByName(editor, languageConfig, cancellationToken) - if (crossSourceFile) { - // TODO (Metrics): 2. Success count for fetchSourceFileByName (find source file by name) - getLogger().debug(`CodeWhisperer finished fetching utg context by file name`) - return { - supplementalContextItems: await generateSupplementalContextFromFocalFile( - crossSourceFile, - 'byName', - cancellationToken - ), - strategy: 'byName', - } - } - throwIfCancelled(cancellationToken) - - crossSourceFile = await findSourceFileByContent(editor, languageConfig, cancellationToken) - if (crossSourceFile) { - // TODO (Metrics): 3. Success count for fetchSourceFileByContent (find source file by content) - getLogger().debug(`CodeWhisperer finished fetching utg context by file content`) - return { - supplementalContextItems: await generateSupplementalContextFromFocalFile( - crossSourceFile, - 'byContent', - cancellationToken - ), - strategy: 'byContent', - } - } - - // TODO (Metrics): 4. Failure count - when unable to find focal file (supplemental context empty) - getLogger().debug(`CodeWhisperer failed to fetch utg context`) - return { - supplementalContextItems: [], - strategy: 'empty', - } -} - -async function generateSupplementalContextFromFocalFile( - filePath: string, - strategy: UtgStrategy, - cancellationToken: vscode.CancellationToken -): Promise { - const fileContent = await fs.readFileText(vscode.Uri.parse(filePath!).fsPath) - - // DO NOT send code chunk with empty content - if (fileContent.trim().length === 0) { - return [] - } - - return [ - { - filePath: filePath, - content: 'UTG\n' + fileContent.slice(0, Math.min(fileContent.length, utgConfig.maxSegmentSize)), - }, - ] -} - -async function findSourceFileByContent( - editor: vscode.TextEditor, - languageConfig: utgLanguageConfig, - cancellationToken: vscode.CancellationToken -): Promise { - const testFileContent = await fs.readFileText(editor.document.fileName) - const testElementList = extractFunctions(testFileContent, languageConfig.functionExtractionPattern) - - throwIfCancelled(cancellationToken) - - testElementList.push(...extractClasses(testFileContent, languageConfig.classExtractionPattern)) - - throwIfCancelled(cancellationToken) - - let sourceFilePath: string | undefined = undefined - let maxMatchCount = 0 - - if (testElementList.length === 0) { - // TODO: Add metrics here, as unable to parse test file using Regex. - return sourceFilePath - } - - const relevantFilePaths = await getRelevantUtgFiles(editor) - - throwIfCancelled(cancellationToken) - - // TODO (Metrics):Add metrics for relevantFilePaths length - for (const filePath of relevantFilePaths) { - throwIfCancelled(cancellationToken) - - const fileContent = await fs.readFileText(filePath) - const elementList = extractFunctions(fileContent, languageConfig.functionExtractionPattern) - elementList.push(...extractClasses(fileContent, languageConfig.classExtractionPattern)) - const matchCount = countSubstringMatches(elementList, testElementList) - if (matchCount > maxMatchCount) { - maxMatchCount = matchCount - sourceFilePath = filePath - } - } - return sourceFilePath -} - -async function getRelevantUtgFiles(editor: vscode.TextEditor): Promise { - const targetFile = editor.document.uri.fsPath - const language = editor.document.languageId - - return await getOpenFilesInWindow(async (candidateFile) => { - return ( - targetFile !== candidateFile && - path.extname(targetFile) === path.extname(candidateFile) && - !(await isTestFile(candidateFile, { languageId: language })) - ) - }) -} - -export function guessSrcFileName( - testFileName: string, - languageId: vscode.TextDocument['languageId'] -): string | undefined { - const languageConfig = utgLanguageConfigs[languageId] - if (!languageConfig) { - return undefined - } - - for (const pattern of languageConfig.testFilenamePattern) { - try { - const match = testFileName.match(pattern) - if (match) { - return match[1] + match[2] - } - } catch (err) { - if (err instanceof Error) { - getLogger().error( - `codewhisperer: error while guessing source file name from file ${testFileName} and pattern ${pattern}: ${err.message}` - ) - } - } - } - - return undefined -} - -async function findSourceFileByName( - editor: vscode.TextEditor, - languageConfig: utgLanguageConfig, - cancellationToken: vscode.CancellationToken -): Promise { - const testFileName = path.basename(editor.document.fileName) - const assumedSrcFileName = guessSrcFileName(testFileName, editor.document.languageId) - if (!assumedSrcFileName) { - return undefined - } - - const sourceFiles = await vscode.workspace.findFiles(`**/${assumedSrcFileName}`) - - throwIfCancelled(cancellationToken) - - if (sourceFiles.length > 0) { - return sourceFiles[0].toString() - } - return undefined -} - -function throwIfCancelled(token: vscode.CancellationToken): void | never { - if (token.isCancellationRequested) { - throw new ToolkitError(supplemetalContextFetchingTimeoutMsg, { cause: new CancellationError('timeout') }) - } -} diff --git a/packages/core/src/codewhisperer/util/telemetryHelper.ts b/packages/core/src/codewhisperer/util/telemetryHelper.ts index 060a5ecb282..bb4e62cfaba 100644 --- a/packages/core/src/codewhisperer/util/telemetryHelper.ts +++ b/packages/core/src/codewhisperer/util/telemetryHelper.ts @@ -101,7 +101,7 @@ export class TelemetryHelper { artifactsUploadDuration: artifactsUploadDuration, buildPayloadBytes: buildPayloadBytes, buildZipFileBytes: buildZipFileBytes, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, acceptedCharactersCount: acceptedCharactersCount, acceptedCount: acceptedCount, acceptedLinesCount: acceptedLinesCount, @@ -146,7 +146,7 @@ export class TelemetryHelper { codewhispererSupplementalContextLength: supplementalContextMetadata?.contentsLength, codewhispererSupplementalContextTimeout: supplementalContextMetadata?.isProcessTimeout, codewhispererTriggerType: session.triggerType, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, duration: duration || 0, reason: reason ? reason.substring(0, 200) : undefined, result, @@ -196,7 +196,7 @@ export class TelemetryHelper { codewhispererTimeToFirstRecommendation: session.timeToFirstRecommendation, codewhispererTriggerType: session.triggerType, codewhispererTypeaheadLength: this.typeAheadLength, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, traceId: this.traceId, }) @@ -283,7 +283,7 @@ export class TelemetryHelper { codewhispererSupplementalContextLength: supplementalContextMetadata?.contentsLength, codewhispererSupplementalContextTimeout: supplementalContextMetadata?.isProcessTimeout, codewhispererTriggerType: session.triggerType, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, traceId: this.traceId, } events.push(event) @@ -658,7 +658,7 @@ export class TelemetryHelper { codewhispererRequestId: this._firstResponseRequestId, codewhispererSessionId: session.sessionId, codewhispererTriggerType: session.triggerType, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, }) } public sendCodeScanEvent(languageId: string, jobId: string) { diff --git a/packages/core/src/codewhisperer/views/activeStateController.ts b/packages/core/src/codewhisperer/views/activeStateController.ts index b3c991a9d38..c1799da8fec 100644 --- a/packages/core/src/codewhisperer/views/activeStateController.ts +++ b/packages/core/src/codewhisperer/views/activeStateController.ts @@ -6,13 +6,9 @@ import * as vscode from 'vscode' import { LineSelection, LinesChangeEvent } from '../tracker/lineTracker' import { isTextEditor } from '../../shared/utilities/editorUtilities' -import { RecommendationService, SuggestionActionEvent } from '../service/recommendationService' import { subscribeOnce } from '../../shared/utilities/vsCodeUtils' import { Container } from '../service/serviceContainer' -import { RecommendationHandler } from '../service/recommendationHandler' import { cancellableDebounce } from '../../shared/utilities/functionUtils' -import { telemetry } from '../../shared/telemetry/telemetry' -import { TelemetryHelper } from '../util/telemetryHelper' export class ActiveStateController implements vscode.Disposable { private readonly _disposable: vscode.Disposable @@ -34,26 +30,21 @@ export class ActiveStateController implements vscode.Disposable { constructor(private readonly container: Container) { this._disposable = vscode.Disposable.from( - RecommendationService.instance.suggestionActionEvent(async (e) => { - await telemetry.withTraceId(async () => { - await this.onSuggestionActionEvent(e) - }, TelemetryHelper.instance.traceId) - }), - RecommendationHandler.instance.onDidReceiveRecommendation(async (_) => { - await this.onDidReceiveRecommendation() - }), + // RecommendationService.instance.suggestionActionEvent(async (e) => { + // await telemetry.withTraceId(async () => { + // await this.onSuggestionActionEvent(e) + // }, TelemetryHelper.instance.traceId) + // }), + // RecommendationHandler.instance.onDidReceiveRecommendation(async (_) => { + // await this.onDidReceiveRecommendation() + // }), this.container.lineTracker.onDidChangeActiveLines(async (e) => { await this.onActiveLinesChanged(e) }), subscribeOnce(this.container.lineTracker.onReady)(async (_) => { await this.onReady() }), - this.container.auth.auth.onDidChangeConnectionState(async (e) => { - if (e.state !== 'authenticating') { - await this._refresh(vscode.window.activeTextEditor) - } - }), - this.container.auth.secondaryAuth.onDidChangeActiveConnection(async () => { + this.container.auth.onDidChangeConnectionState(async (e) => { await this._refresh(vscode.window.activeTextEditor) }) ) @@ -70,32 +61,32 @@ export class ActiveStateController implements vscode.Disposable { await this._refresh(vscode.window.activeTextEditor) } - private async onSuggestionActionEvent(e: SuggestionActionEvent) { - if (!this._isReady) { - return - } - - this.clear(e.editor) // do we need this? - if (e.triggerType === 'OnDemand' && e.isRunning) { - // if user triggers on demand, immediately update the UI and cancel the previous debounced update if there is one - this.refreshDebounced.cancel() - await this._refresh(this._editor) - } else { - await this.refreshDebounced.promise(e.editor) - } - } - - private async onDidReceiveRecommendation() { - if (!this._isReady) { - return - } - - if (this._editor && this._editor === vscode.window.activeTextEditor) { - // receives recommendation, immediately update the UI and cacnel the debounced update if there is one - this.refreshDebounced.cancel() - await this._refresh(this._editor, false) - } - } + // private async onSuggestionActionEvent(e: SuggestionActionEvent) { + // if (!this._isReady) { + // return + // } + + // this.clear(e.editor) // do we need this? + // if (e.triggerType === 'OnDemand' && e.isRunning) { + // // if user triggers on demand, immediately update the UI and cancel the previous debounced update if there is one + // this.refreshDebounced.cancel() + // await this._refresh(this._editor) + // } else { + // await this.refreshDebounced.promise(e.editor) + // } + // } + + // private async onDidReceiveRecommendation() { + // if (!this._isReady) { + // return + // } + + // if (this._editor && this._editor === vscode.window.activeTextEditor) { + // // receives recommendation, immediately update the UI and cacnel the debounced update if there is one + // this.refreshDebounced.cancel() + // await this._refresh(this._editor, false) + // } + // } private async onActiveLinesChanged(e: LinesChangeEvent) { if (!this._isReady) { @@ -139,7 +130,7 @@ export class ActiveStateController implements vscode.Disposable { return } - if (!this.container.auth.isConnectionValid()) { + if (!this.container.auth.isConnected()) { this.clear(this._editor) return } @@ -147,7 +138,7 @@ export class ActiveStateController implements vscode.Disposable { if (shouldDisplay !== undefined) { await this.updateDecorations(editor, selections, shouldDisplay) } else { - await this.updateDecorations(editor, selections, RecommendationService.instance.isRunning) + await this.updateDecorations(editor, selections, true) } } diff --git a/packages/core/src/codewhisperer/views/lineAnnotationController.ts b/packages/core/src/codewhisperer/views/lineAnnotationController.ts index 8b1d38ed7ae..3d1055769c7 100644 --- a/packages/core/src/codewhisperer/views/lineAnnotationController.ts +++ b/packages/core/src/codewhisperer/views/lineAnnotationController.ts @@ -9,18 +9,16 @@ import { LineSelection, LinesChangeEvent } from '../tracker/lineTracker' import { isTextEditor } from '../../shared/utilities/editorUtilities' import { cancellableDebounce } from '../../shared/utilities/functionUtils' import { subscribeOnce } from '../../shared/utilities/vsCodeUtils' -import { RecommendationService } from '../service/recommendationService' +// import { RecommendationService } from '../service/recommendationService' import { AnnotationChangeSource, inlinehintKey } from '../models/constants' import globals from '../../shared/extensionGlobals' import { Container } from '../service/serviceContainer' import { telemetry } from '../../shared/telemetry/telemetry' import { getLogger } from '../../shared/logger/logger' -import { Commands } from '../../shared/vscode/commands2' -import { session } from '../util/codeWhispererSession' -import { RecommendationHandler } from '../service/recommendationHandler' +// import { session } from '../util/codeWhispererSession' +// import { RecommendationHandler } from '../service/recommendationHandler' import { runtimeLanguageContext } from '../util/runtimeLanguageContext' import { setContext } from '../../shared/vscode/setContext' -import { TelemetryHelper } from '../util/telemetryHelper' const case3TimeWindow = 30000 // 30 seconds @@ -75,13 +73,14 @@ export class AutotriggerState implements AnnotationState { static acceptedCount = 0 updateState(changeSource: AnnotationChangeSource, force: boolean): AnnotationState | undefined { - if (AutotriggerState.acceptedCount < RecommendationService.instance.acceptedSuggestionCount) { - return new ManualtriggerState() - } else if (session.recommendations.length > 0 && RecommendationHandler.instance.isSuggestionVisible()) { - return new PressTabState() - } else { - return this - } + // if (AutotriggerState.acceptedCount < RecommendationService.instance.acceptedSuggestionCount) { + // return new ManualtriggerState() + // } else if (session.recommendations.length > 0 && RecommendationHandler.instance.isSuggestionVisible()) { + // return new PressTabState() + // } else { + // return this + // } + return undefined } isNextState(state: AnnotationState | undefined): boolean { @@ -268,50 +267,45 @@ export class LineAnnotationController implements vscode.Disposable { subscribeOnce(this.container.lineTracker.onReady)(async (_) => { await this.onReady() }), - RecommendationService.instance.suggestionActionEvent(async (e) => { - await telemetry.withTraceId(async () => { - if (!this._isReady) { - return - } - - if (this._currentState instanceof ManualtriggerState) { - if (e.triggerType === 'OnDemand' && this._currentState.hasManualTrigger === false) { - this._currentState.hasManualTrigger = true - } - if ( - e.response?.recommendationCount !== undefined && - e.response?.recommendationCount > 0 && - this._currentState.hasValidResponse === false - ) { - this._currentState.hasValidResponse = true - } - } - - await this.refresh(e.editor, 'codewhisperer') - }, TelemetryHelper.instance.traceId) - }), + // RecommendationService.instance.suggestionActionEvent(async (e) => { + // await telemetry.withTraceId(async () => { + // if (!this._isReady) { + // return + // } + + // if (this._currentState instanceof ManualtriggerState) { + // if (e.triggerType === 'OnDemand' && this._currentState.hasManualTrigger === false) { + // this._currentState.hasManualTrigger = true + // } + // if ( + // e.response?.recommendationCount !== undefined && + // e.response?.recommendationCount > 0 && + // this._currentState.hasValidResponse === false + // ) { + // this._currentState.hasValidResponse = true + // } + // } + + // await this.refresh(e.editor, 'codewhisperer') + // }, TelemetryHelper.instance.traceId) + // }), this.container.lineTracker.onDidChangeActiveLines(async (e) => { await this.onActiveLinesChanged(e) }), - this.container.auth.auth.onDidChangeConnectionState(async (e) => { - if (e.state !== 'authenticating') { - await this.refresh(vscode.window.activeTextEditor, 'editor') - } - }), - this.container.auth.secondaryAuth.onDidChangeActiveConnection(async () => { + this.container.auth.onDidChangeConnectionState(async () => { await this.refresh(vscode.window.activeTextEditor, 'editor') - }), - Commands.register('aws.amazonq.dismissTutorial', async () => { - const editor = vscode.window.activeTextEditor - if (editor) { - this.clear() - try { - telemetry.ui_click.emit({ elementId: `dismiss_${this._currentState.id}` }) - } catch (_) {} - await this.dismissTutorial() - getLogger().debug(`codewhisperer: user dismiss tutorial.`) - } }) + // Commands.register('aws.amazonq.dismissTutorial', async () => { + // const editor = vscode.window.activeTextEditor + // if (editor) { + // this.clear() + // try { + // telemetry.ui_click.emit({ elementId: `dismiss_${this._currentState.id}` }) + // } catch (_) {} + // await this.dismissTutorial() + // getLogger().debug(`codewhisperer: user dismiss tutorial.`) + // } + // }) ) } @@ -424,7 +418,7 @@ export class LineAnnotationController implements vscode.Disposable { return } - if (!this.container.auth.isConnectionValid()) { + if (!this.container.auth.isConnected()) { this.clear() return } @@ -484,7 +478,7 @@ export class LineAnnotationController implements vscode.Disposable { source: AnnotationChangeSource, force?: boolean ): Partial | undefined { - const isCWRunning = RecommendationService.instance.isRunning + const isCWRunning = true const textOptions: vscode.ThemableDecorationAttachmentRenderOptions = { contentText: '', @@ -517,9 +511,9 @@ export class LineAnnotationController implements vscode.Disposable { this._currentState = updatedState // take snapshot of accepted session so that we can compre if there is delta -> users accept 1 suggestion after seeing this state - AutotriggerState.acceptedCount = RecommendationService.instance.acceptedSuggestionCount + AutotriggerState.acceptedCount = 0 // take snapshot of total trigger count so that we can compare if there is delta -> users accept/reject suggestions after seeing this state - TryMoreExState.triggerCount = RecommendationService.instance.totalValidTriggerCount + TryMoreExState.triggerCount = 0 textOptions.contentText = this._currentState.text() diff --git a/packages/core/src/codewhisperer/views/securityIssue/securityIssueWebview.ts b/packages/core/src/codewhisperer/views/securityIssue/securityIssueWebview.ts index d511bd9a5f6..0426c853d9b 100644 --- a/packages/core/src/codewhisperer/views/securityIssue/securityIssueWebview.ts +++ b/packages/core/src/codewhisperer/views/securityIssue/securityIssueWebview.ts @@ -135,7 +135,7 @@ export class SecurityIssueWebview extends VueWebview { ruleId: this.issue!.ruleId, component: 'webview', result: 'Succeeded', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, codeFixAction: fixAction, } } diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts index ba2072eb6dc..eaa35e6087f 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/controller.ts @@ -59,19 +59,15 @@ import { triggerPayloadToChatRequest } from './chatRequest/converter' import { AuthUtil } from '../../../codewhisperer/util/authUtil' import { openUrl } from '../../../shared/utilities/vsCodeUtils' import { randomUUID } from '../../../shared/crypto' -import { LspController } from '../../../amazonq/lsp/lspController' import { CodeWhispererSettings } from '../../../codewhisperer/util/codewhispererSettings' import { getSelectedCustomization } from '../../../codewhisperer/util/customizationUtil' import { getHttpStatusCode, AwsClientResponseError } from '../../../shared/errors' import { uiEventRecorder } from '../../../amazonq/util/eventRecorder' import { telemetry } from '../../../shared/telemetry/telemetry' -import { isSsoConnection } from '../../../auth/connection' import { inspect } from '../../../shared/utilities/collectionUtils' import { DefaultAmazonQAppInitContext } from '../../../amazonq/apps/initContext' import globals from '../../../shared/extensionGlobals' import { MynahIconsType, MynahUIDataModel, QuickActionCommand } from '@aws/mynah-ui' -import { LspClient } from '../../../amazonq/lsp/lspClient' -import { AdditionalContextPrompt, ContextCommandItem, ContextCommandItemType } from '../../../amazonq/lsp/types' import { workspaceCommand } from '../../../amazonq/webview/ui/tabs/constants' import fs from '../../../shared/fs/fs' import { FeatureConfigProvider, Features } from '../../../shared/featureConfig' @@ -80,9 +76,6 @@ import { getUserPromptsDirectory, promptFileExtension, createSavedPromptCommandId, - aditionalContentNameLimit, - additionalContentInnerContextLimit, - workspaceChunkMaxSize, defaultContextLengths, } from '../../constants' import { ChatSession } from '../../clients/chat/v0/chat' @@ -527,7 +520,7 @@ export class ChatController { commands: [{ command: commandName, description: commandDescription }], }) } - const symbolsCmd: QuickActionCommand = contextCommand[0].commands?.[3] + // const symbolsCmd: QuickActionCommand = contextCommand[0].commands?.[3] const promptsCmd: QuickActionCommand = contextCommand[0].commands?.[4] // Check for user prompts @@ -543,7 +536,7 @@ export class ChatController { command: path.basename(name, promptFileExtension), icon: 'magic' as MynahIconsType, id: 'prompt', - label: 'file' as ContextCommandItemType, + // label: 'file' as ContextCommandItemType, route: [userPromptsDirectory, name], })) ) @@ -559,47 +552,47 @@ export class ChatController { icon: 'list-add' as MynahIconsType, }) - const lspClientReady = await LspClient.instance.waitUntilReady() - if (lspClientReady) { - const contextCommandItems = await LspClient.instance.getContextCommandItems() - const folderCmd: QuickActionCommand = contextCommand[0].commands?.[1] - const filesCmd: QuickActionCommand = contextCommand[0].commands?.[2] - - for (const contextCommandItem of contextCommandItems) { - const wsFolderName = path.basename(contextCommandItem.workspaceFolder) - if (contextCommandItem.type === 'file') { - filesCmd.children?.[0].commands.push({ - command: path.basename(contextCommandItem.relativePath), - description: path.join(wsFolderName, contextCommandItem.relativePath), - route: [contextCommandItem.workspaceFolder, contextCommandItem.relativePath], - label: 'file' as ContextCommandItemType, - id: contextCommandItem.id, - icon: 'file' as MynahIconsType, - }) - } else if (contextCommandItem.type === 'folder') { - folderCmd.children?.[0].commands.push({ - command: path.basename(contextCommandItem.relativePath), - description: path.join(wsFolderName, contextCommandItem.relativePath), - route: [contextCommandItem.workspaceFolder, contextCommandItem.relativePath], - label: 'folder' as ContextCommandItemType, - id: contextCommandItem.id, - icon: 'folder' as MynahIconsType, - }) - } else if (contextCommandItem.symbol && symbolsCmd.children) { - symbolsCmd.children?.[0].commands.push({ - command: contextCommandItem.symbol.name, - description: `${contextCommandItem.symbol.kind}, ${path.join(wsFolderName, contextCommandItem.relativePath)}, L${contextCommandItem.symbol.range.start.line}-${contextCommandItem.symbol.range.end.line}`, - route: [contextCommandItem.workspaceFolder, contextCommandItem.relativePath], - label: 'code' as ContextCommandItemType, - id: contextCommandItem.id, - icon: 'code-block' as MynahIconsType, - }) - } - } - } + // const lspClientReady = await LspClient.instance.waitUntilReady() + // if (lspClientReady) { + // const contextCommandItems = await LspClient.instance.getContextCommandItems() + // const folderCmd: QuickActionCommand = contextCommand[0].commands?.[1] + // const filesCmd: QuickActionCommand = contextCommand[0].commands?.[2] + + // for (const contextCommandItem of contextCommandItems) { + // const wsFolderName = path.basename(contextCommandItem.workspaceFolder) + // if (contextCommandItem.type === 'file') { + // filesCmd.children?.[0].commands.push({ + // command: path.basename(contextCommandItem.relativePath), + // description: path.join(wsFolderName, contextCommandItem.relativePath), + // route: [contextCommandItem.workspaceFolder, contextCommandItem.relativePath], + // label: 'file' as ContextCommandItemType, + // id: contextCommandItem.id, + // icon: 'file' as MynahIconsType, + // }) + // } else if (contextCommandItem.type === 'folder') { + // folderCmd.children?.[0].commands.push({ + // command: path.basename(contextCommandItem.relativePath), + // description: path.join(wsFolderName, contextCommandItem.relativePath), + // route: [contextCommandItem.workspaceFolder, contextCommandItem.relativePath], + // label: 'folder' as ContextCommandItemType, + // id: contextCommandItem.id, + // icon: 'folder' as MynahIconsType, + // }) + // } else if (contextCommandItem.symbol && symbolsCmd.children) { + // symbolsCmd.children?.[0].commands.push({ + // command: contextCommandItem.symbol.name, + // description: `${contextCommandItem.symbol.kind}, ${path.join(wsFolderName, contextCommandItem.relativePath)}, L${contextCommandItem.symbol.range.start.line}-${contextCommandItem.symbol.range.end.line}`, + // route: [contextCommandItem.workspaceFolder, contextCommandItem.relativePath], + // label: 'code' as ContextCommandItemType, + // id: contextCommandItem.id, + // icon: 'code-block' as MynahIconsType, + // }) + // } + // } + // } this.messenger.sendContextCommandData(contextCommand) - void LspController.instance.updateContextCommandSymbolsOnce() + // void LspController.instance.updateContextCommandSymbolsOnce() } private handlePromptCreate(tabID: string) { @@ -967,9 +960,9 @@ export class ChatController { const tabID = triggerEvent.tabID - const credentialsState = await AuthUtil.instance.getChatAuthState() + const credentialsState = AuthUtil.instance.getAuthState() - if (credentialsState.codewhispererChat !== 'connected' && credentialsState.codewhispererCore !== 'connected') { + if (credentialsState !== 'connected') { await this.messenger.sendAuthNeededExceptionMessage(credentialsState, tabID, triggerID) return } @@ -1006,7 +999,7 @@ export class ChatController { } private async resolveContextCommandPayload(triggerPayload: TriggerPayload, session: ChatSession) { - const contextCommands: ContextCommandItem[] = [] + const contextCommands: any[] = [] // Check for workspace rules to add to context const workspaceRules = await this.collectWorkspaceRules() @@ -1017,7 +1010,7 @@ export class ChatController { vscode.workspace.getWorkspaceFolder(vscode.Uri.parse(rule))?.uri?.path || '' return { workspaceFolder: workspaceFolderPath, - type: 'file' as ContextCommandItemType, + type: 'file' as any, relativePath: path.relative(workspaceFolderPath, rule), } }) @@ -1029,7 +1022,7 @@ export class ChatController { if (typeof context !== 'string' && context.route && context.route.length === 2) { contextCommands.push({ workspaceFolder: context.route[0] || '', - type: (context.label || '') as ContextCommandItemType, + type: (context.label || '') as any, relativePath: context.route[1] || '', id: context.id, }) @@ -1044,45 +1037,45 @@ export class ChatController { return [] } workspaceFolders.sort() - const workspaceFolder = workspaceFolders[0] - for (const contextCommand of contextCommands) { - session.relativePathToWorkspaceRoot.set(contextCommand.workspaceFolder, contextCommand.workspaceFolder) - } - let prompts: AdditionalContextPrompt[] = [] - try { - prompts = await LspClient.instance.getContextCommandPrompt(contextCommands) - } catch (e) { - // todo: handle @workspace used before indexing is ready - getLogger().verbose(`Could not get context command prompts: ${e}`) - } - - triggerPayload.contextLengths.additionalContextLengths = this.telemetryHelper.getContextLengths(prompts) - for (const prompt of prompts.slice(0, 20)) { - // Add system prompt for user prompts and workspace rules - const contextType = this.telemetryHelper.getContextType(prompt) - const description = - contextType === 'rule' || contextType === 'prompt' - ? `You must follow the instructions in ${prompt.relativePath}. Below are lines ${prompt.startLine}-${prompt.endLine} of this file:\n` - : prompt.description - - // Handle user prompts outside the workspace - const relativePath = prompt.filePath.startsWith(getUserPromptsDirectory()) - ? path.basename(prompt.filePath) - : path.relative(workspaceFolder, prompt.filePath) - - const entry = { - name: prompt.name.substring(0, aditionalContentNameLimit), - description: description.substring(0, aditionalContentNameLimit), - innerContext: prompt.content.substring(0, additionalContentInnerContextLimit), - type: contextType, - relativePath: relativePath, - startLine: prompt.startLine, - endLine: prompt.endLine, - } - - triggerPayload.additionalContents.push(entry) - } - getLogger().info(`Retrieved chunks of additional context count: ${triggerPayload.additionalContents.length} `) + // const workspaceFolder = workspaceFolders[0] + // for (const contextCommand of contextCommands) { + // session.relativePathToWorkspaceRoot.set(contextCommand.workspaceFolder, contextCommand.workspaceFolder) + // } + // const prompts: any[] = [] + // try { + // // prompts = await LspClient.instance.getContextCommandPrompt(contextCommands) + // } catch (e) { + // // todo: handle @workspace used before indexing is ready + // getLogger().verbose(`Could not get context command prompts: ${e}`) + // } + + // triggerPayload.contextLengths.additionalContextLengths = this.telemetryHelper.getContextLengths(prompts) + // for (const prompt of prompts.slice(0, 20)) { + // // Add system prompt for user prompts and workspace rules + // const contextType = this.telemetryHelper.getContextType(prompt) + // const description = + // contextType === 'rule' || contextType === 'prompt' + // ? `You must follow the instructions in ${prompt.relativePath}. Below are lines ${prompt.startLine}-${prompt.endLine} of this file:\n` + // : prompt.description + + // // Handle user prompts outside the workspace + // const relativePath = prompt.filePath.startsWith(getUserPromptsDirectory()) + // ? path.basename(prompt.filePath) + // : path.relative(workspaceFolder, prompt.filePath) + + // const entry = { + // name: prompt.name.substring(0, aditionalContentNameLimit), + // description: description.substring(0, aditionalContentNameLimit), + // innerContext: prompt.content.substring(0, additionalContentInnerContextLimit), + // type: contextType, + // relativePath: relativePath, + // startLine: prompt.startLine, + // endLine: prompt.endLine, + // } + + // triggerPayload.additionalContents.push(entry) + // } + // getLogger().info(`Retrieved chunks of additional context count: ${triggerPayload.additionalContents.length} `) } private async generateResponse( @@ -1109,11 +1102,9 @@ export class ChatController { const tabID = triggerEvent.tabID - const credentialsState = await AuthUtil.instance.getChatAuthState() + const credentialsState = AuthUtil.instance.getAuthState() - if ( - !(credentialsState.codewhispererChat === 'connected' && credentialsState.codewhispererCore === 'connected') - ) { + if (!(credentialsState === 'connected')) { await this.messenger.sendAuthNeededExceptionMessage(credentialsState, tabID, triggerID) return } @@ -1130,25 +1121,24 @@ export class ChatController { if (triggerPayload.useRelevantDocuments) { triggerPayload.message = triggerPayload.message.replace(/@workspace/, '') if (CodeWhispererSettings.instance.isLocalIndexEnabled()) { - const start = performance.now() - const relevantTextDocuments = await LspController.instance.query(triggerPayload.message) - for (const relevantDocument of relevantTextDocuments) { - if (relevantDocument.text && relevantDocument.text.length > 0) { - triggerPayload.contextLengths.workspaceContextLength += relevantDocument.text.length - if (relevantDocument.text.length > workspaceChunkMaxSize) { - relevantDocument.text = relevantDocument.text.substring(0, workspaceChunkMaxSize) - getLogger().debug(`Truncating @workspace chunk: ${relevantDocument.relativeFilePath} `) - } - triggerPayload.relevantTextDocuments.push(relevantDocument) - } - } - - for (const doc of triggerPayload.relevantTextDocuments) { - getLogger().info( - `amazonq: Using workspace files ${doc.relativeFilePath}, content(partial): ${doc.text?.substring(0, 200)}, start line: ${doc.startLine}, end line: ${doc.endLine}` - ) - } - triggerPayload.projectContextQueryLatencyMs = performance.now() - start + // const start = performance.now() + // const relevantTextDocuments = await LspController.instance.query(triggerPayload.message) + // for (const relevantDocument of relevantTextDocuments) { + // if (relevantDocument.text && relevantDocument.text.length > 0) { + // triggerPayload.contextLengths.workspaceContextLength += relevantDocument.text.length + // if (relevantDocument.text.length > workspaceChunkMaxSize) { + // relevantDocument.text = relevantDocument.text.substring(0, workspaceChunkMaxSize) + // getLogger().debug(`Truncating @workspace chunk: ${relevantDocument.relativeFilePath} `) + // } + // triggerPayload.relevantTextDocuments.push(relevantDocument) + // } + // } + // for (const doc of triggerPayload.relevantTextDocuments) { + // getLogger().info( + // `amazonq: Using workspace files ${doc.relativeFilePath}, content(partial): ${doc.text?.substring(0, 200)}, start line: ${doc.startLine}, end line: ${doc.endLine}` + // ) + // } + // triggerPayload.projectContextQueryLatencyMs = performance.now() - start } else { this.messenger.sendOpenSettingsMessage(triggerID, tabID) return @@ -1193,7 +1183,7 @@ export class ChatController { try { this.messenger.sendInitalStream(tabID, triggerID, triggerPayload.documentReferences) this.telemetryHelper.setConversationStreamStartTime(tabID) - if (isSsoConnection(AuthUtil.instance.conn)) { + if (AuthUtil.instance.isConnected() && AuthUtil.instance.isSsoSession()) { const { $metadata, generateAssistantResponseResponse } = await session.chatSso(request) response = { $metadata: $metadata, diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts index 8c914686ad4..20a05fcbdcf 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts @@ -36,17 +36,16 @@ import { ChatPromptCommandType, DocumentReference, TriggerPayload } from '../mod import { getHttpStatusCode, getRequestId, ToolkitError } from '../../../../shared/errors' import { keys } from '../../../../shared/utilities/tsUtils' import { getLogger } from '../../../../shared/logger/logger' -import { FeatureAuthState } from '../../../../codewhisperer/util/authUtil' import { CodeScanIssue } from '../../../../codewhisperer/models/model' import { marked } from 'marked' import { JSDOM } from 'jsdom' -import { LspController } from '../../../../amazonq/lsp/lspController' import { extractCodeBlockLanguage } from '../../../../shared/markdown' import { extractAuthFollowUp } from '../../../../amazonq/util/authUtils' import { helpMessage } from '../../../../amazonq/webview/ui/texts/constants' import { ChatItem, ChatItemButton, ChatItemFormItem, DetailedList, MynahUIDataModel } from '@aws/mynah-ui' import { Database } from '../../../../shared/db/chatDb/chatDb' import { TabType } from '../../../../amazonq/webview/ui/storages/tabsStorage' +import { AuthState } from '../../../../auth/auth2' export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help' @@ -63,7 +62,7 @@ export class Messenger { private readonly telemetryHelper: CWCTelemetryHelper ) {} - public async sendAuthNeededExceptionMessage(credentialState: FeatureAuthState, tabID: string, triggerID: string) { + public async sendAuthNeededExceptionMessage(credentialState: AuthState, tabID: string, triggerID: string) { const { message, authType } = extractAuthFollowUp(credentialState) this.dispatcher.sendAuthNeededExceptionMessage( new AuthNeededException( @@ -290,11 +289,7 @@ export class Messenger { relatedContent: { title: 'Sources', content: relatedSuggestions as any }, }) } - if ( - triggerPayload.relevantTextDocuments && - triggerPayload.relevantTextDocuments.length > 0 && - LspController.instance.isIndexingInProgress() - ) { + if (triggerPayload.relevantTextDocuments && triggerPayload.relevantTextDocuments.length > 0) { this.dispatcher.sendChatMessage( new ChatMessage( { diff --git a/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts b/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts index 2d9e01db9a0..10720959e13 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts @@ -2,7 +2,7 @@ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. * SPDX-License-Identifier: Apache-2.0 */ -import * as path from 'path' + import { UserIntent } from '@amzn/codewhisperer-streaming' import { AmazonqAddMessage, @@ -28,7 +28,6 @@ import { ResponseBodyLinkClickMessage, SourceLinkClickMessage, TriggerPayload, - AdditionalContextLengths, AdditionalContextInfo, } from './model' import { TriggerEvent, TriggerEventsStorage } from '../../storages/triggerEvents' @@ -43,9 +42,6 @@ import { supportedLanguagesList } from '../chat/chatRequest/converter' import { AuthUtil } from '../../../codewhisperer/util/authUtil' import { getSelectedCustomization } from '../../../codewhisperer/util/customizationUtil' import { undefinedIfEmpty } from '../../../shared/utilities/textUtilities' -import { AdditionalContextPrompt } from '../../../amazonq/lsp/types' -import { getUserPromptsDirectory, promptFileExtension } from '../../constants' -import { isInDirectory } from '../../../shared/filesystemUtilities' import { sleep } from '../../../shared/utilities/timeoutUtils' import { FileDiagnostic, @@ -71,7 +67,7 @@ export function recordTelemetryChatRunCommand(type: CwsprChatCommandType, comman result: 'Succeeded', cwsprChatCommandType: type, cwsprChatCommandName: command, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, }) } @@ -164,40 +160,6 @@ export class CWCTelemetryHelper { telemetry.amazonq_exitFocusChat.emit({ result: 'Succeeded', passive: true }) } - public getContextType(prompt: AdditionalContextPrompt): string { - if (prompt.filePath.endsWith(promptFileExtension)) { - if (isInDirectory(path.join('.amazonq', 'rules'), prompt.relativePath)) { - return 'rule' - } else if (isInDirectory(getUserPromptsDirectory(), prompt.filePath)) { - return 'prompt' - } - } - return 'file' - } - - public getContextLengths(prompts: AdditionalContextPrompt[]): AdditionalContextLengths { - let fileContextLength = 0 - let promptContextLength = 0 - let ruleContextLength = 0 - - for (const prompt of prompts) { - const type = this.getContextType(prompt) - switch (type) { - case 'rule': - ruleContextLength += prompt.content.length - break - case 'file': - fileContextLength += prompt.content.length - break - case 'prompt': - promptContextLength += prompt.content.length - break - } - } - - return { fileContextLength, promptContextLength, ruleContextLength } - } - public async recordFeedback(message: ChatItemFeedbackMessage) { const logger = getLogger() try { @@ -255,7 +217,7 @@ export class CWCTelemetryHelper { event = { result, cwsprChatConversationId: conversationId ?? '', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, cwsprChatMessageId: message.messageId, cwsprChatUserIntent: this.getUserIntentForTelemetry(message.userIntent), cwsprChatInteractionType: 'insertAtCursor', @@ -273,7 +235,7 @@ export class CWCTelemetryHelper { event = { result, cwsprChatConversationId: conversationId ?? '', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, cwsprChatMessageId: message.messageId, cwsprChatUserIntent: this.getUserIntentForTelemetry(message.userIntent), cwsprChatInteractionType: 'copySnippet', @@ -292,7 +254,7 @@ export class CWCTelemetryHelper { cwsprChatConversationId: conversationId ?? '', cwsprChatMessageId: message.messageId, cwsprChatInteractionType: 'acceptDiff', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, cwsprChatAcceptedCharactersLength: message.code.length, cwsprChatHasReference: message.referenceTrackerInformation && message.referenceTrackerInformation.length > 0, @@ -307,7 +269,7 @@ export class CWCTelemetryHelper { cwsprChatConversationId: conversationId ?? '', cwsprChatMessageId: message.messageId, cwsprChatInteractionType: 'viewDiff', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, cwsprChatAcceptedCharactersLength: message.code.length, cwsprChatHasReference: message.referenceTrackerInformation && message.referenceTrackerInformation.length > 0, @@ -320,7 +282,7 @@ export class CWCTelemetryHelper { event = { result, cwsprChatConversationId: conversationId ?? '', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, cwsprChatMessageId: message.messageId, cwsprChatInteractionType: 'clickFollowUp', } @@ -331,7 +293,7 @@ export class CWCTelemetryHelper { result, cwsprChatMessageId: message.messageId, cwsprChatConversationId: conversationId ?? '', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, cwsprChatInteractionType: message.vote, } break @@ -341,7 +303,7 @@ export class CWCTelemetryHelper { result, cwsprChatMessageId: message.messageId, cwsprChatConversationId: conversationId ?? '', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, cwsprChatInteractionType: 'clickLink', cwsprChatInteractionTarget: message.link, } @@ -352,7 +314,7 @@ export class CWCTelemetryHelper { result, cwsprChatMessageId: message.messageId, cwsprChatConversationId: conversationId ?? '', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, cwsprChatInteractionType: 'clickBodyLink', cwsprChatInteractionTarget: message.link, } @@ -363,7 +325,7 @@ export class CWCTelemetryHelper { result, cwsprChatMessageId: 'footer', cwsprChatConversationId: conversationId ?? '', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, cwsprChatInteractionType: 'clickBodyLink', cwsprChatInteractionTarget: message.link, } @@ -479,7 +441,7 @@ export class CWCTelemetryHelper { cwsprChatUserIntent: telemetryUserIntent, cwsprChatHasCodeSnippet: triggerPayload.codeSelection && !triggerPayload.codeSelection.isEmpty, cwsprChatProgrammingLanguage: triggerPayload.fileLanguage, - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, cwsprChatHasProjectContext: triggerPayload.relevantTextDocuments ? triggerPayload.relevantTextDocuments.length > 0 && triggerPayload.useRelevantDocuments === true : false, @@ -567,7 +529,7 @@ export class CWCTelemetryHelper { cwsprChatRequestLength: triggerPayload.message.length, cwsprChatResponseLength: message.messageLength, cwsprChatConversationType: 'Chat', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, codewhispererCustomizationArn: triggerPayload.customization.arn, cwsprChatHasProjectContext: hasProjectLevelContext, cwsprChatHasContextList: triggerPayload.documentReferences.length > 0, @@ -641,7 +603,7 @@ export class CWCTelemetryHelper { cwsprChatResponseCode: responseCode, cwsprChatRequestLength: triggerPayload.message?.length ?? 0, cwsprChatConversationType: 'Chat', - credentialStartUrl: AuthUtil.instance.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, }) } diff --git a/packages/core/src/codewhispererChat/editor/codelens.ts b/packages/core/src/codewhispererChat/editor/codelens.ts index 4df72d776d6..4e671315fc4 100644 --- a/packages/core/src/codewhispererChat/editor/codelens.ts +++ b/packages/core/src/codewhispererChat/editor/codelens.ts @@ -8,7 +8,7 @@ import { ToolkitError } from '../../shared/errors' import { Commands, placeholder } from '../../shared/vscode/commands2' import { platform } from 'os' import { focusAmazonQPanel } from '../commands/registerCommands' -import { AuthStates, AuthUtil } from '../../codewhisperer/util/authUtil' +import { AuthUtil } from '../../codewhisperer/util/authUtil' import { inlinehintKey } from '../../codewhisperer/models/constants' import { EndState } from '../../codewhisperer/views/lineAnnotationController' @@ -89,7 +89,7 @@ export class TryChatCodeLensProvider implements vscode.CodeLensProvider { return resolve([]) } - if (AuthUtil.instance.getChatAuthStateSync().amazonQ !== AuthStates.connected) { + if (AuthUtil.instance.getAuthState() !== 'connected') { return resolve([]) } diff --git a/packages/core/src/dev/activation.ts b/packages/core/src/dev/activation.ts index 8ce0f6aab11..16b5d7e53ad 100644 --- a/packages/core/src/dev/activation.ts +++ b/packages/core/src/dev/activation.ts @@ -25,7 +25,6 @@ import { NotificationsController } from '../notifications/controller' import { DevNotificationsState } from '../notifications/types' import { QuickPickItem } from 'vscode' import { ChildProcess } from '../shared/utilities/processUtils' -import { WorkspaceLspInstaller } from '../amazonq/lsp/workspaceInstaller' interface MenuOption { readonly label: string @@ -451,12 +450,6 @@ const resettableFeatures: readonly ResettableFeature[] = [ detail: 'Resets memory/global state for the notifications panel (includes dismissed, onReceive).', executor: resetNotificationsState, }, - { - name: 'workspace lsp', - label: 'Download Lsp ', - detail: 'Resets workspace LSP', - executor: resetWorkspaceLspDownload, - }, ] as const // TODO this is *somewhat* similar to `openStorageFromInput`. If we need another @@ -545,10 +538,6 @@ async function resetNotificationsState() { await targetNotificationsController.reset() } -async function resetWorkspaceLspDownload() { - await new WorkspaceLspInstaller().resolve() -} - async function editNotifications() { const storageKey = 'aws.notifications.dev' const current = globalState.get(storageKey) ?? {} diff --git a/packages/core/src/login/webview/commonAuthViewProvider.ts b/packages/core/src/login/webview/commonAuthViewProvider.ts index f805d7cf759..639267c5fc0 100644 --- a/packages/core/src/login/webview/commonAuthViewProvider.ts +++ b/packages/core/src/login/webview/commonAuthViewProvider.ts @@ -44,9 +44,8 @@ import { CodeCatalystAuthenticationProvider } from '../../codecatalyst/auth' import { telemetry } from '../../shared/telemetry/telemetry' import { AuthSources } from './util' import { AuthFlowStates } from './vue/types' -import { getTelemetryMetadataForConn } from '../../auth/connection' -import { AuthUtil } from '../../codewhisperer/util/authUtil' import { ExtensionUse } from '../../auth/utils' +import { AuthUtil } from '../../codewhisperer/util/authUtil' export class CommonAuthViewProvider implements WebviewViewProvider { public readonly viewType: string @@ -109,7 +108,7 @@ export class CommonAuthViewProvider implements WebviewViewProvider { if (authState === AuthFlowStates.REAUTHNEEDED || authState === AuthFlowStates.REAUTHENTICATING) { this.webView!.server.storeMetricMetadata({ isReAuth: true, - ...(await getTelemetryMetadataForConn(AuthUtil.instance.conn)), + ...(await AuthUtil.instance.getTelemetryMetadata()), }) } else { this.webView!.server.storeMetricMetadata({ isReAuth: false }) diff --git a/packages/core/src/login/webview/index.ts b/packages/core/src/login/webview/index.ts index 80abcc4fd79..d7760c89d4c 100644 --- a/packages/core/src/login/webview/index.ts +++ b/packages/core/src/login/webview/index.ts @@ -5,3 +5,4 @@ export { CommonAuthViewProvider } from './commonAuthViewProvider' export { CommonAuthWebview } from './vue/backend' +export * as backendAmazonQ from './vue/amazonq/backend_amazonq' diff --git a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts index a83e99d04b7..0a9dd576d6f 100644 --- a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts +++ b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts @@ -3,13 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ import * as vscode from 'vscode' -import { - AwsConnection, - Connection, - SsoConnection, - getTelemetryMetadataForConn, - isSsoConnection, -} from '../../../../auth/connection' +import { AwsConnection, SsoConnection } from '../../../../auth/connection' import { AuthUtil } from '../../../../codewhisperer/util/authUtil' import { CommonAuthWebview } from '../backend' import { awsIdSignIn } from '../../../../codewhisperer/util/showSsoPrompt' @@ -70,9 +64,8 @@ export class AmazonQLoginWebview extends CommonAuthWebview { authEnabledFeatures: 'codewhisperer', isReAuth: false, }) - - const conn = await awsIdSignIn() - this.storeMetricMetadata(await getTelemetryMetadataForConn(conn)) + await awsIdSignIn() + this.storeMetricMetadata(await AuthUtil.instance.getTelemetryMetadata()) void vscode.window.showInformationMessage('AmazonQ: Successfully connected to AWS Builder ID') }) @@ -92,8 +85,8 @@ export class AmazonQLoginWebview extends CommonAuthWebview { isReAuth: false, }) - const conn = await connectToEnterpriseSso(startUrl, region) - this.storeMetricMetadata(await getTelemetryMetadataForConn(conn)) + await connectToEnterpriseSso(startUrl, region) + this.storeMetricMetadata(await AuthUtil.instance.getTelemetryMetadata()) void vscode.window.showInformationMessage('AmazonQ: Successfully connected to AWS IAM Identity Center') }) @@ -109,9 +102,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview { getLogger().error('amazon Q reauthenticate called on a non-existant connection') throw new ToolkitError('Cannot reauthenticate non-existant connection.') } - - const conn = AuthUtil.instance.conn - if (!isSsoConnection(conn)) { + if (!AuthUtil.instance.isSsoSession()) { getLogger().error('amazon Q reauthenticate called, but the connection is not SSO') throw new ToolkitError('Cannot reauthenticate non-SSO connection.') } @@ -122,14 +113,12 @@ export class AmazonQLoginWebview extends CommonAuthWebview { */ this.reauthError = await this.ssoSetup('reauthenticateAmazonQ', async () => { this.storeMetricMetadata({ - authEnabledFeatures: this.getAuthEnabledFeatures(conn), + authEnabledFeatures: 'codewhisperer', isReAuth: true, - ...(await getTelemetryMetadataForConn(conn)), + ...(await AuthUtil.instance.getTelemetryMetadata()), }) await AuthUtil.instance.reauthenticate() - this.storeMetricMetadata({ - ...(await getTelemetryMetadataForConn(conn)), - }) + this.storeMetricMetadata(await AuthUtil.instance.getTelemetryMetadata()) }) } finally { this.isReauthenticating = false @@ -147,10 +136,6 @@ export class AmazonQLoginWebview extends CommonAuthWebview { return this.reauthError } - async getActiveConnection(): Promise { - return AuthUtil.instance.conn - } - /** * `true` if the actual reauth flow is in progress. * @@ -162,11 +147,13 @@ export class AmazonQLoginWebview extends CommonAuthWebview { isReauthenticating: boolean = false private authState: AuthFlowState = 'LOGIN' override async refreshAuthState(): Promise { - const featureAuthStates = await AuthUtil.instance.getChatAuthState() - if (featureAuthStates.amazonQ === 'expired') { + if (AuthUtil.instance.getAuthState() === 'expired') { this.authState = this.isReauthenticating ? 'REAUTHENTICATING' : 'REAUTHNEEDED' return - } else if (featureAuthStates.amazonQ === 'pendingProfileSelection') { + } else if ( + AuthUtil.instance.isConnected() && + AuthUtil.instance.regionProfileManager.requireProfileSelection() + ) { this.authState = 'PENDING_PROFILE_SELECTION' // possible that user starts with "profile selection" state therefore the timeout for auth flow should be disposed otherwise will emit failure this.loadMetadata?.loadTimeout?.dispose() @@ -186,19 +173,18 @@ export class AmazonQLoginWebview extends CommonAuthWebview { @withTelemetryContext({ name: 'signout', class: className }) override async signout(): Promise { - const conn = AuthUtil.instance.secondaryAuth.activeConnection - if (!isSsoConnection(conn)) { - throw new ToolkitError(`Cannot signout non-SSO connection, type is: ${conn?.type}`) + if (!AuthUtil.instance.isSsoSession()) { + throw new ToolkitError(`Cannot signout non-SSO connection`) } this.storeMetricMetadata({ - authEnabledFeatures: this.getAuthEnabledFeatures(conn), + authEnabledFeatures: 'codewhisperer', isReAuth: true, - ...(await getTelemetryMetadataForConn(conn)), + ...(await AuthUtil.instance.getTelemetryMetadata()), result: 'Cancelled', }) - await AuthUtil.instance.secondaryAuth.deleteConnection() + await AuthUtil.instance.logout() this.reauthError = undefined this.emitAuthMetric() @@ -229,13 +215,12 @@ export class AmazonQLoginWebview extends CommonAuthWebview { try { return await AuthUtil.instance.regionProfileManager.getProfiles() } catch (e) { - const conn = AuthUtil.instance.conn as SsoConnection | undefined telemetry.amazonq_didSelectProfile.emit({ source: 'auth', amazonQProfileRegion: AuthUtil.instance.regionProfileManager.activeRegionProfile?.region ?? 'not-set', - ssoRegion: conn?.ssoRegion, + ssoRegion: AuthUtil.instance.connection?.region, result: 'Failed', - credentialStartUrl: conn?.startUrl, + credentialStartUrl: AuthUtil.instance.connection?.startUrl, reason: (e as Error).message, }) @@ -250,7 +235,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview { private setupConnectionEventEmitter(): void { // allows the frontend to listen to Amazon Q auth events from the backend const codeWhispererConnectionChanged = createThrottle(() => this.onActiveConnectionModified.fire()) - AuthUtil.instance.secondaryAuth.onDidChangeActiveConnection(codeWhispererConnectionChanged) + AuthUtil.instance.onDidChangeConnectionState(codeWhispererConnectionChanged) AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(codeWhispererConnectionChanged) /** diff --git a/packages/core/src/login/webview/vue/backend.ts b/packages/core/src/login/webview/vue/backend.ts index c8f1f38d4d7..edb1980a8c0 100644 --- a/packages/core/src/login/webview/vue/backend.ts +++ b/packages/core/src/login/webview/vue/backend.ts @@ -14,7 +14,6 @@ import { handleWebviewError } from '../../../webviews/server' import { InvalidGrantException } from '@aws-sdk/client-sso-oidc' import { AwsConnection, - Connection, hasScopes, scopesCodeCatalyst, scopesCodeWhispererChat, @@ -196,8 +195,6 @@ export abstract class CommonAuthWebview extends VueWebview { abstract reauthenticateConnection(): Promise abstract getReauthError(): Promise - abstract getActiveConnection(): Promise - /** Refreshes the current state of the auth flow, determining what you see in the UI */ abstract refreshAuthState(): Promise /** Use {@link refreshAuthState} first to ensure this returns the latest state */ diff --git a/packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts b/packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts index 4e4db35b9ad..caec2c764bc 100644 --- a/packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts +++ b/packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts @@ -9,7 +9,6 @@ import { getLogger } from '../../../../shared/logger/logger' import { CommonAuthWebview } from '../backend' import { AwsConnection, - Connection, SsoConnection, TelemetryMetadata, createSsoProfile, @@ -161,9 +160,6 @@ export class ToolkitLoginWebview extends CommonAuthWebview { override reauthenticateConnection(): Promise { throw new Error('Method not implemented.') } - override getActiveConnection(): Promise { - throw new Error('Method not implemented.') - } override async refreshAuthState(): Promise {} override async getAuthState(): Promise { diff --git a/packages/core/src/shared/clients/codewhispererChatClient.ts b/packages/core/src/shared/clients/codewhispererChatClient.ts index 93b7b3bceb4..1210adb7e30 100644 --- a/packages/core/src/shared/clients/codewhispererChatClient.ts +++ b/packages/core/src/shared/clients/codewhispererChatClient.ts @@ -10,7 +10,7 @@ import { getUserAgent } from '../telemetry/util' // Create a client for featureDev streaming based off of aws sdk v3 export async function createCodeWhispererChatStreamingClient(): Promise { - const bearerToken = await AuthUtil.instance.getBearerToken() + const bearerToken = await AuthUtil.instance.getToken() const cwsprConfig = getCodewhispererConfig() const streamingClient = new CodeWhispererStreaming({ region: cwsprConfig.region, diff --git a/packages/core/src/shared/clients/qDeveloperChatClient.ts b/packages/core/src/shared/clients/qDeveloperChatClient.ts index ee98a78e356..d9344b5b406 100644 --- a/packages/core/src/shared/clients/qDeveloperChatClient.ts +++ b/packages/core/src/shared/clients/qDeveloperChatClient.ts @@ -6,12 +6,13 @@ import { QDeveloperStreaming } from '@amzn/amazon-q-developer-streaming-client' import { getCodewhispererConfig } from '../../codewhisperer/client/codewhisperer' import { getUserAgent } from '../telemetry/util' import { ConfiguredRetryStrategy } from '@smithy/util-retry' -import { AuthUtil } from '../../codewhisperer/util/authUtil' // Create a client for featureDev streaming based off of aws sdk v3 export async function createQDeveloperStreamingClient(): Promise { + throw new Error('Do not call this function until IAM is supported by LSP identity server') + const cwsprConfig = getCodewhispererConfig() - const credentials = await AuthUtil.instance.getCredentials() + const credentials = undefined const streamingClient = new QDeveloperStreaming({ region: cwsprConfig.region, endpoint: cwsprConfig.endpoint, diff --git a/packages/core/src/shared/featureConfig.ts b/packages/core/src/shared/featureConfig.ts index d7acb9657be..7cc6a9cbfc7 100644 --- a/packages/core/src/shared/featureConfig.ts +++ b/packages/core/src/shared/featureConfig.ts @@ -13,7 +13,6 @@ import * as nls from 'vscode-nls' import { codeWhispererClient as client } from '../codewhisperer/client/codewhisperer' import { AuthUtil } from '../codewhisperer/util/authUtil' import { getLogger } from './logger/logger' -import { isBuilderIdConnection, isIdcSsoConnection } from '../auth/connection' import { CodeWhispererSettings } from '../codewhisperer/util/codewhispererSettings' import globals from './extensionGlobals' import { getClientId, getOperatingSystem } from './telemetry/util' @@ -149,9 +148,9 @@ export class FeatureConfigProvider { const previousOverride = globals.globalState.tryGet('aws.amazonq.customization.overrideV2', String) if (customizationArnOverride !== undefined && customizationArnOverride !== previousOverride) { // Double check if server-side wrongly returns a customizationArn to BID users - if (isBuilderIdConnection(AuthUtil.instance.conn)) { + if (AuthUtil.instance.isBuilderIdConnection()) { this.featureConfigs.delete(Features.customizationArnOverride) - } else if (isIdcSsoConnection(AuthUtil.instance.conn)) { + } else if (AuthUtil.instance.isIdcConnection()) { const availableCustomizations = await getAvailableCustomizationsList() // If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value diff --git a/packages/core/src/shared/globalState.ts b/packages/core/src/shared/globalState.ts index 13db46b430a..65d761412b8 100644 --- a/packages/core/src/shared/globalState.ts +++ b/packages/core/src/shared/globalState.ts @@ -7,6 +7,7 @@ import * as vscode from 'vscode' import { getLogger } from './logger/logger' import * as redshift from '../awsService/redshift/models/models' import { TypeConstructor, cast } from './utilities/typeConstructors' +import { Customization } from '../codewhisperer/client/codewhispereruserclient' type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId' export type ToolIdStateKey = `${ToolId}.savedConnectionId` @@ -229,6 +230,53 @@ export class GlobalState implements vscode.Memento { return all?.[warehouseArn] } + /** + * Get the Amazon Q customization. If legacy (map of customizations) store the + * customization with label of profile name + * + * @param profileName name of profile, only used in case legacy customization is found + * @returns Amazon Q customization, or undefined if not found. + * If legacy, return the Amazon Q customization for the auth profile name + */ + getAmazonQCustomization(profileName: string): Customization | undefined { + const result = this.tryGet('CODEWHISPERER_SELECTED_CUSTOMIZATION', Object, undefined) + + // Legacy migration for old customization map of type { [label: string]: Customization[] } + if (typeof result === 'object' && Object.values(result).every(Array.isArray)) { + const selectedCustomization = result[profileName] + this.tryUpdate('CODEWHISPERER_SELECTED_CUSTOMIZATION', selectedCustomization) + return selectedCustomization + } else { + return result + } + } + + /** + * Get the Amazon Q cached customizations. If legacy (map of customizations) store the + * customizations with label of profile name + * + * @param profileName name of profile, only used in case legacy customization is found + * @returns array of Amazon Q cached customizations, or empty array if not found. + * If legacy, return the Amazon Q persisted customizations for the auth profile name + */ + getAmazonQCachedCustomization(profileName: string): Customization[] { + const result = this.tryGet('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', Array, []) + + // Legacy migration for old customization map of type { [label: string]: Customization[] } + if (result.length === 0) { + const customizations = this.tryGet<{ [label: string]: Customization[] }>( + 'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', + Object, + {} + ) + const cachedCustomizationsArray = customizations[profileName] || [] + this.tryUpdate('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', cachedCustomizationsArray) + return cachedCustomizationsArray + } else { + return result + } + } + /** * Sets SSO session creation timestamp for the given session `id`. * @@ -270,3 +318,68 @@ export class GlobalState implements vscode.Memento { return all?.[id] } } + +export interface GlobalStatePollerProps { + getState: () => any + changeHandler: () => void + pollIntervalInMs: number +} + +/** + * Utility class that polls a state value at regular intervals and triggers a callback when the state changes. + * + * This class can be used to monitor changes in global state and react to those changes. + */ +export class GlobalStatePoller { + protected oldValue: any + protected pollIntervalInMs: number + protected getState: () => any + protected changeHandler: () => void + protected intervalId?: NodeJS.Timeout + + constructor(props: GlobalStatePollerProps) { + this.getState = props.getState + this.changeHandler = props.changeHandler + this.pollIntervalInMs = props.pollIntervalInMs + this.oldValue = this.getState() + } + + /** + * Factory method that creates and starts a GlobalStatePoller instance. + * + * @param getState - Function that returns the current state value to monitor, e.g. globals.globalState.tryGet + * @param changeHandler - Callback function that is invoked when the state changes + * @returns A new GlobalStatePoller instance that has already started polling + */ + static create(props: GlobalStatePollerProps) { + const instance = new GlobalStatePoller(props) + instance.poll() + return instance + } + + /** + * Starts polling the state value. When a change is detected, the changeHandler callback is invoked. + */ + private poll() { + if (this.intervalId) { + this.kill() + } + this.intervalId = setInterval(() => { + const newValue = this.getState() + if (this.oldValue !== newValue) { + this.oldValue = newValue + this.changeHandler() + } + }, this.pollIntervalInMs) + } + + /** + * Stops the polling interval. + */ + kill() { + if (this.intervalId) { + clearInterval(this.intervalId) + this.intervalId = undefined + } + } +} diff --git a/packages/core/src/shared/index.ts b/packages/core/src/shared/index.ts index f4c78e2093c..b74b6862b45 100644 --- a/packages/core/src/shared/index.ts +++ b/packages/core/src/shared/index.ts @@ -49,6 +49,9 @@ export * as env from './vscode/env' export * from './vscode/commands2' export * from './utilities/pathUtils' export * from './utilities/zipStream' +export * as editorUtilities from './utilities/editorUtilities' +export * as functionUtilities from './utilities/functionUtils' +export * as vscodeUtilities from './utilities/vsCodeUtils' export * from './errors' export * as messages from './utilities/messages' export * as errors from './errors' @@ -74,3 +77,4 @@ export * as BaseLspInstaller from './lsp/baseLspInstaller' export * as collectionUtil from './utilities/collectionUtils' export * from './datetime' export * from './performance/marks' +export * as mementoUtils from './utilities/mementos' diff --git a/packages/core/src/shared/logger/logger.ts b/packages/core/src/shared/logger/logger.ts index b398ff93162..bb94fb0dc53 100644 --- a/packages/core/src/shared/logger/logger.ts +++ b/packages/core/src/shared/logger/logger.ts @@ -21,6 +21,7 @@ export type LogTopic = | 'nextEditPrediction' | 'resourceCache' | 'telemetry' + | 'amazonqAuth' class ErrorLog { constructor( diff --git a/packages/core/src/shared/lsp/baseLspInstaller.ts b/packages/core/src/shared/lsp/baseLspInstaller.ts index 0aeca1dfda4..7acf58ad788 100644 --- a/packages/core/src/shared/lsp/baseLspInstaller.ts +++ b/packages/core/src/shared/lsp/baseLspInstaller.ts @@ -5,7 +5,6 @@ import * as nodePath from 'path' import vscode from 'vscode' -import { LspConfig } from '../../amazonq/lsp/config' import { LanguageServerResolver } from './lspResolver' import { ManifestResolver } from './manifestResolver' import { LspResolution, ResourcePaths } from './types' @@ -14,6 +13,14 @@ import { Range } from 'semver' import { getLogger } from '../logger/logger' import type { Logger, LogTopic } from '../logger/logger' +export interface LspConfig { + manifestUrl: string + supportedVersions: string + id: string + suppressPromptPrefix: string + path?: string +} + export abstract class BaseLspInstaller { private logger: Logger diff --git a/packages/core/src/shared/settings-toolkit.gen.ts b/packages/core/src/shared/settings-toolkit.gen.ts index 59a637a4870..10020cf51f9 100644 --- a/packages/core/src/shared/settings-toolkit.gen.ts +++ b/packages/core/src/shared/settings-toolkit.gen.ts @@ -44,6 +44,7 @@ export const toolkitSettings = { "jsonResourceModification": {}, "amazonqLSP": {}, "amazonqLSPInline": {}, + "amazonqLSPInlineChat": {}, "amazonqChatLSP": {} }, "aws.resources.enabledResources": {}, diff --git a/packages/core/src/shared/telemetry/exemptMetrics.ts b/packages/core/src/shared/telemetry/exemptMetrics.ts index a3fc8d5ad78..4e0deacc058 100644 --- a/packages/core/src/shared/telemetry/exemptMetrics.ts +++ b/packages/core/src/shared/telemetry/exemptMetrics.ts @@ -29,6 +29,8 @@ const validationExemptMetrics: Set = new Set([ 'codewhisperer_codePercentage', 'codewhisperer_userModification', 'codewhisperer_userTriggerDecision', + 'codewhisperer_perceivedLatency', // flare doesn't currently set result property + 'codewhisperer_serviceInvocation', // flare doesn't currently set result property 'dynamicresource_selectResources', 'dynamicresource_copyIdentifier', 'dynamicresource_mutateResource', diff --git a/packages/core/src/shared/utilities/functionUtils.ts b/packages/core/src/shared/utilities/functionUtils.ts index cbf89340ade..214721b1cdb 100644 --- a/packages/core/src/shared/utilities/functionUtils.ts +++ b/packages/core/src/shared/utilities/functionUtils.ts @@ -93,9 +93,10 @@ export function memoize(fn: (...args: U) => T): (...args: U) */ export function debounce( cb: (...args: Input) => Output | Promise, - delay: number = 0 + delay: number = 0, + useLastCall: boolean = false ): (...args: Input) => Promise { - return cancellableDebounce(cb, delay).promise + return cancellableDebounce(cb, delay, useLastCall).promise } /** @@ -104,10 +105,12 @@ export function debounce( */ export function cancellableDebounce( cb: (...args: Input) => Output | Promise, - delay: number = 0 + delay: number = 0, + useLastCall: boolean = false ): { promise: (...args: Input) => Promise; cancel: () => void } { let timeout: Timeout | undefined let promise: Promise | undefined + let lastestArgs: Input | undefined const cancel = (): void => { if (timeout) { @@ -119,6 +122,7 @@ export function cancellableDebounce( return { promise: (...args: Input) => { + lastestArgs = args timeout?.refresh() return (promise ??= new Promise((resolve, reject) => { @@ -126,7 +130,8 @@ export function cancellableDebounce( timeout.onCompletion(async () => { timeout = promise = undefined try { - resolve(await cb(...args)) + const argsToUse = useLastCall ? lastestArgs! : args + resolve(await cb(...argsToUse)) } catch (err) { reject(err) } diff --git a/packages/core/src/shared/utilities/index.ts b/packages/core/src/shared/utilities/index.ts index 520390b5204..ecf753090ca 100644 --- a/packages/core/src/shared/utilities/index.ts +++ b/packages/core/src/shared/utilities/index.ts @@ -6,3 +6,4 @@ export { isExtensionInstalled, isExtensionActive } from './vsCodeUtils' export { VSCODE_EXTENSION_ID } from '../extensions' export * from './functionUtils' +export * as messageUtils from './messages' diff --git a/packages/core/src/test/amazonq/customizationUtil.test.ts b/packages/core/src/test/amazonq/customizationUtil.test.ts index a3a49e907d9..19e59b91c03 100644 --- a/packages/core/src/test/amazonq/customizationUtil.test.ts +++ b/packages/core/src/test/amazonq/customizationUtil.test.ts @@ -21,22 +21,17 @@ import { import { FeatureContext, globals } from '../../shared' import { resetCodeWhispererGlobalVariables } from '../codewhisperer/testUtil' import { createSsoProfile, createTestAuth } from '../credentials/testUtil' -import { SsoConnection } from '../../auth' +import { createTestAuthUtil } from '../testAuthUtil' const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start' describe('customizationProvider', function () { - let auth: ReturnType - let ssoConn: SsoConnection let regionProfileManager: RegionProfileManager beforeEach(async () => { - auth = createTestAuth(globals.globalState) - ssoConn = await auth.createInvalidSsoConnection( - createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: amazonQScopes }) - ) - - regionProfileManager = new RegionProfileManager(() => ssoConn) + createTestAuth(globals.globalState) + await createTestAuthUtil() + regionProfileManager = new RegionProfileManager(AuthUtil.instance) }) afterEach(() => { @@ -66,7 +61,6 @@ describe('customizationProvider', function () { describe('CodeWhisperer-customizationUtils', function () { let auth: ReturnType - let ssoConn: SsoConnection let featureCustomization: FeatureContext before(async function () { @@ -75,8 +69,10 @@ describe('CodeWhisperer-customizationUtils', function () { }) beforeEach(async function () { + await createTestAuthUtil() + auth = createTestAuth(globals.globalState) - ssoConn = await auth.createInvalidSsoConnection( + await auth.createInvalidSsoConnection( createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: amazonQScopes }) ) featureCustomization = { @@ -91,7 +87,6 @@ describe('CodeWhisperer-customizationUtils', function () { sinon.stub(AuthUtil.instance, 'isConnectionExpired').returns(false) sinon.stub(AuthUtil.instance, 'isConnected').returns(true) sinon.stub(AuthUtil.instance, 'isCustomizationFeatureEnabled').value(true) - sinon.stub(AuthUtil.instance, 'conn').value(ssoConn) await resetCodeWhispererGlobalVariables() }) @@ -101,14 +96,15 @@ describe('CodeWhisperer-customizationUtils', function () { }) it('Returns baseCustomization when not SSO', async function () { - sinon.stub(AuthUtil.instance, 'isValidEnterpriseSsoInUse').returns(false) + sinon.stub(AuthUtil.instance, 'isIdcConnection').returns(false) + const customization = getSelectedCustomization() assert.strictEqual(customization.name, baseCustomization.name) }) it('Returns selectedCustomization when customization manually selected', async function () { - sinon.stub(AuthUtil.instance, 'isValidEnterpriseSsoInUse').returns(true) + sinon.stub(AuthUtil.instance, 'isIdcConnection').returns(true) const selectedCustomization: Customization = { arn: 'selectedCustomizationArn', @@ -124,6 +120,8 @@ describe('CodeWhisperer-customizationUtils', function () { }) it(`setSelectedCustomization should set to the customization provided if override option is false or not specified`, async function () { + sinon.stub(AuthUtil.instance, 'isIdcConnection').returns(true) + await setSelectedCustomization({ arn: 'FOO' }, false) assert.strictEqual(getSelectedCustomization().arn, 'FOO') @@ -138,6 +136,8 @@ describe('CodeWhisperer-customizationUtils', function () { }) it(`setSelectedCustomization should only set to the customization provided once for override per customization arn if override is true`, async function () { + sinon.stub(AuthUtil.instance, 'isIdcConnection').returns(true) + await setSelectedCustomization({ arn: 'OVERRIDE' }, true) assert.strictEqual(getSelectedCustomization().arn, 'OVERRIDE') diff --git a/packages/core/src/test/amazonqDoc/controller.test.ts b/packages/core/src/test/amazonqDoc/controller.test.ts index d69edc47fd7..e3c62ffb2b8 100644 --- a/packages/core/src/test/amazonqDoc/controller.test.ts +++ b/packages/core/src/test/amazonqDoc/controller.test.ts @@ -18,7 +18,6 @@ import { } from './utils' import { CurrentWsFolders, MetricDataOperationName, MetricDataResult, NewFileInfo } from '../../amazonqDoc/types' import { DocCodeGenState, docScheme, Session } from '../../amazonqDoc' -import { AuthUtil } from '../../codewhisperer' import { ApiClientError, ApiServiceError, @@ -49,6 +48,7 @@ import { WorkspaceEmptyError, } from '../../amazonqDoc/errors' import { LlmError } from '../../amazonq/errors' + describe('Controller - Doc Generation', () => { const firstTabID = '123' const firstConversationID = '123' @@ -176,11 +176,6 @@ describe('Controller - Doc Generation', () => { configurable: true, }) - sandbox.stub(AuthUtil.instance, 'getChatAuthState').resolves({ - codewhispererCore: 'connected', - codewhispererChat: 'connected', - amazonQ: 'connected', - }) sandbox.stub(FileSystem.prototype, 'exists').resolves(false) if (isMultiTabs) { const secondSession = await createCodeGenState(sandbox, secondTabID, secondConversationID, secondUploadID) diff --git a/packages/core/src/test/amazonqDoc/utils.ts b/packages/core/src/test/amazonqDoc/utils.ts index 51c7305902c..d6d74e7ac3c 100644 --- a/packages/core/src/test/amazonqDoc/utils.ts +++ b/packages/core/src/test/amazonqDoc/utils.ts @@ -24,6 +24,8 @@ import { MetricData, } from '../../amazonqFeatureDev/client/featuredevproxyclient' import { FollowUpTypes } from '../../amazonq/commons/types' +import { AuthUtil } from '../../codewhisperer/util/authUtil' +import { LanguageClientAuth } from '../../auth/auth2' export function createMessenger(sandbox: sinon.SinonSandbox): DocMessenger { return new DocMessenger( @@ -102,7 +104,17 @@ export async function sessionWriteFile(session: Session, uri: vscode.Uri, encode }) } +export function createMockAuthUtil(sandbox: sinon.SinonSandbox) { + const mockLspAuth: Partial = { + registerSsoTokenChangedHandler: sinon.stub().resolves(), + } + AuthUtil.create(mockLspAuth as LanguageClientAuth) + sandbox.stub(AuthUtil.instance.regionProfileManager, 'onDidChangeRegionProfile').resolves() + sandbox.stub(AuthUtil.instance, 'getAuthState').returns('connected') +} + export async function createController(sandbox: sinon.SinonSandbox): Promise { + createMockAuthUtil(sandbox) const messenger = createMessenger(sandbox) // Create a new workspace root diff --git a/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts b/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts index 7848d0561b0..8e8639332fd 100644 --- a/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts +++ b/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts @@ -141,11 +141,7 @@ describe('Controller', () => { scheme: featureDevScheme, }) - sinon.stub(AuthUtil.instance, 'getChatAuthState').resolves({ - codewhispererCore: 'connected', - codewhispererChat: 'connected', - amazonQ: 'connected', - }) + sinon.stub(AuthUtil.instance, 'getAuthState').returns('connected') }) afterEach(() => { diff --git a/packages/core/src/test/codewhisperer/commands/basicCommands.test.ts b/packages/core/src/test/codewhisperer/commands/basicCommands.test.ts index 01c7c43c947..7fb7107af47 100644 --- a/packages/core/src/test/codewhisperer/commands/basicCommands.test.ts +++ b/packages/core/src/test/codewhisperer/commands/basicCommands.test.ts @@ -56,7 +56,7 @@ import { waitUntil } from '../../../shared/utilities/timeoutUtils' import { listCodeWhispererCommands } from '../../../codewhisperer/ui/statusBarMenu' import { CodeScanIssue, CodeScansState, CodeSuggestionsState, codeScanState } from '../../../codewhisperer/models/model' import { cwQuickPickSource } from '../../../codewhisperer/commands/types' -import { refreshStatusBar } from '../../../codewhisperer/service/inlineCompletionService' +import { refreshStatusBar } from '../../../codewhisperer/service/statusBar' import { focusAmazonQPanel } from '../../../codewhispererChat/commands/registerCommands' import * as diagnosticsProvider from '../../../codewhisperer/service/diagnosticsProvider' import { randomUUID } from '../../../shared/crypto' @@ -409,8 +409,8 @@ describe('CodeWhisperer-basicCommands', function () { it('includes the "source" in the command execution metric', async function () { tryRegister(focusAmazonQPanel) - sinon.stub(AuthUtil.instance.secondaryAuth, 'deleteConnection') - targetCommand = testCommand(signoutCodeWhisperer, AuthUtil.instance) + sinon.stub(AuthUtil.instance, 'logout') + targetCommand = testCommand(signoutCodeWhisperer) await targetCommand.execute(placeholder, cwQuickPickSource) assertTelemetry('vscode_executeCommand', [ { source: cwQuickPickSource, command: focusAmazonQPanel.id }, @@ -475,8 +475,9 @@ describe('CodeWhisperer-basicCommands', function () { it('also shows customizations when connected to valid sso', async function () { sinon.stub(AuthUtil.instance, 'isConnectionExpired').returns(false) sinon.stub(AuthUtil.instance, 'isConnected').returns(true) - sinon.stub(AuthUtil.instance, 'isValidEnterpriseSsoInUse').returns(true) + sinon.stub(AuthUtil.instance, 'isIdcConnection').returns(true) sinon.stub(AuthUtil.instance, 'isCustomizationFeatureEnabled').value(true) + sinon.stub(AuthUtil.instance.regionProfileManager, 'requireProfileSelection').returns(false) await CodeScansState.instance.setScansEnabled(false) getTestWindow().onDidShowQuickPick(async (e) => { @@ -499,7 +500,7 @@ describe('CodeWhisperer-basicCommands', function () { it('should not show auto-scans if using builder id', async function () { sinon.stub(AuthUtil.instance, 'isConnected').returns(true) - sinon.stub(AuthUtil.instance, 'isBuilderIdInUse').returns(true) + sinon.stub(AuthUtil.instance, 'isBuilderIdConnection').returns(true) getTestWindow().onDidShowQuickPick(async (e) => { e.assertItems([ diff --git a/packages/core/src/test/codewhisperer/customizationUtil.test.ts b/packages/core/src/test/codewhisperer/customizationUtil.test.ts new file mode 100644 index 00000000000..ab81a50bafb --- /dev/null +++ b/packages/core/src/test/codewhisperer/customizationUtil.test.ts @@ -0,0 +1,47 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as sinon from 'sinon' +import * as assert from 'assert' +import * as customizationModule from '../../../src/codewhisperer/util/customizationUtil' + +describe('getNewCustomizations', () => { + let getPersistedCustomizationsStub: sinon.SinonStub + + const availableCustomizations = [ + { arn: 'arn1', name: 'custom1' }, + { arn: 'arn2', name: 'custom2' }, + ] + + const persistedCustomizations = [[{ arn: 'arn1', name: 'custom1' }], [{ arn: 'arn2', name: 'custom2' }]] + + beforeEach(() => { + getPersistedCustomizationsStub = sinon.stub(customizationModule, 'getPersistedCustomizations') + }) + + afterEach(() => { + sinon.restore() + }) + + it('returns new customizations that are not in persisted customizations', () => { + const customizations = [...availableCustomizations, { arn: 'arn3', name: 'custom3' }] + + getPersistedCustomizationsStub.returns(persistedCustomizations) + + const result = customizationModule.getNewCustomizations(customizations) + + assert.deepEqual(result, [{ arn: 'arn3', name: 'custom3' }]) + sinon.assert.calledOnce(getPersistedCustomizationsStub) + }) + + it('returns empty array when all available customizations are persisted', () => { + getPersistedCustomizationsStub.returns(persistedCustomizations) + + const result = customizationModule.getNewCustomizations(availableCustomizations) + + assert.deepEqual(result.length, 0) + sinon.assert.calledOnce(getPersistedCustomizationsStub) + }) +}) diff --git a/packages/core/src/test/codewhisperer/startSecurityScan.test.ts b/packages/core/src/test/codewhisperer/startSecurityScan.test.ts index 551949aa3ab..38b00a2bdd3 100644 --- a/packages/core/src/test/codewhisperer/startSecurityScan.test.ts +++ b/packages/core/src/test/codewhisperer/startSecurityScan.test.ts @@ -30,6 +30,7 @@ import * as errors from '../../shared/errors' import * as timeoutUtils from '../../shared/utilities/timeoutUtils' import { SecurityIssueTreeViewProvider } from '../../codewhisperer' import { createClient, mockGetCodeScanResponse } from './testUtil' +import { createTestAuthUtil } from '../testAuthUtil' let extensionContext: FakeExtensionContext let mockSecurityPanelViewProvider: SecurityPanelViewProvider @@ -40,7 +41,9 @@ let focusStub: sinon.SinonStub describe('startSecurityScan', function () { const workspaceFolder = getTestWorkspaceFolder() + beforeEach(async function () { + await createTestAuthUtil() extensionContext = await FakeExtensionContext.create() mockSecurityPanelViewProvider = new SecurityPanelViewProvider(extensionContext) appRoot = join(workspaceFolder, 'python3.7-plain-sam-app') @@ -50,9 +53,11 @@ describe('startSecurityScan', function () { sinon.stub(timeoutUtils, 'sleep') focusStub = sinon.stub(SecurityIssueTreeViewProvider, 'focus') }) + afterEach(function () { sinon.restore() }) + after(async function () { await closeAllEditors() }) diff --git a/packages/core/src/test/codewhisperer/testUtil.ts b/packages/core/src/test/codewhisperer/testUtil.ts index f3b82fd3850..dd8188b1006 100644 --- a/packages/core/src/test/codewhisperer/testUtil.ts +++ b/packages/core/src/test/codewhisperer/testUtil.ts @@ -14,7 +14,6 @@ import { } from '../../codewhisperer/models/model' import { MockDocument } from '../fake/fakeDocument' import { getLogger } from '../../shared/logger' -import { CodeWhispererCodeCoverageTracker } from '../../codewhisperer/tracker/codewhispererCodeCoverageTracker' import globals from '../../shared/extensionGlobals' import { session } from '../../codewhisperer/util/codeWhispererSession' import { DefaultAWSClientBuilder, ServiceOptions } from '../../shared/awsClientBuilder' @@ -23,7 +22,6 @@ import { HttpResponse, Service } from 'aws-sdk' import userApiConfig = require('./../../codewhisperer/client/user-service-2.json') import CodeWhispererUserClient = require('../../codewhisperer/client/codewhispereruserclient') import { codeWhispererClient } from '../../codewhisperer/client/codewhisperer' -import { RecommendationHandler } from '../../codewhisperer/service/recommendationHandler' import * as model from '../../codewhisperer/models/model' import { stub } from '../utilities/stubber' import { Dirent } from 'fs' // eslint-disable-line no-restricted-imports @@ -31,12 +29,10 @@ import { Dirent } from 'fs' // eslint-disable-line no-restricted-imports export async function resetCodeWhispererGlobalVariables() { vsCodeState.isIntelliSenseActive = false vsCodeState.isCodeWhispererEditing = false - CodeWhispererCodeCoverageTracker.instances.clear() globals.telemetry.logger.clear() session.reset() await globals.globalState.clear() await CodeSuggestionsState.instance.setSuggestionsEnabled(true) - await RecommendationHandler.instance.clearInlineCompletionStates() } export function createMockDocument( diff --git a/packages/core/src/test/codewhisperer/zipUtil.test.ts b/packages/core/src/test/codewhisperer/zipUtil.test.ts index a82db4a6840..e6c4f4148e5 100644 --- a/packages/core/src/test/codewhisperer/zipUtil.test.ts +++ b/packages/core/src/test/codewhisperer/zipUtil.test.ts @@ -16,7 +16,6 @@ import { ToolkitError } from '../../shared/errors' import { fs } from '../../shared/fs/fs' import { tempDirPath } from '../../shared/filesystemUtilities' import { CodeWhispererConstants } from '../../codewhisperer/indexNode' -import { LspClient } from '../../amazonq/lsp/lspClient' describe('zipUtil', function () { const workspaceFolder = getTestWorkspaceFolder() @@ -179,23 +178,5 @@ describe('zipUtil', function () { assert.strictEqual(result.language, 'java') assert.strictEqual(result.scannedFiles.size, 4) }) - - it('Should handle file system errors during directory creation', async function () { - sinon.stub(LspClient, 'instance').get(() => ({ - getRepoMapJSON: sinon.stub().resolves('{"mock": "data"}'), - })) - sinon.stub(fs, 'mkdir').rejects(new Error('Directory creation failed')) - - await assert.rejects(() => zipUtil.generateZipTestGen(appRoot, false), /Directory creation failed/) - }) - - it('Should handle zip project errors', async function () { - sinon.stub(LspClient, 'instance').get(() => ({ - getRepoMapJSON: sinon.stub().resolves('{"mock": "data"}'), - })) - sinon.stub(zipUtil, 'zipProject' as keyof ZipUtil).rejects(new Error('Zip failed')) - - await assert.rejects(() => zipUtil.generateZipTestGen(appRoot, false), /Zip failed/) - }) }) }) diff --git a/packages/core/src/test/codewhispererChat/editor/codelens.test.ts b/packages/core/src/test/codewhispererChat/editor/codelens.test.ts index 52243027ebf..3e9bd9284a5 100644 --- a/packages/core/src/test/codewhispererChat/editor/codelens.test.ts +++ b/packages/core/src/test/codewhispererChat/editor/codelens.test.ts @@ -15,7 +15,7 @@ import { InstalledClock } from '@sinonjs/fake-timers' import globals from '../../../shared/extensionGlobals' import { focusAmazonQPanel } from '../../../codewhispererChat/commands/registerCommands' import sinon from 'sinon' -import { AuthState, AuthStates, AuthUtil, FeatureAuthState } from '../../../codewhisperer/util/authUtil' +import { AuthUtil } from '../../../codewhisperer/util/authUtil' import { inlinehintKey } from '../../../codewhisperer/models/constants' import { AutotriggerState, @@ -24,6 +24,8 @@ import { PressTabState, TryMoreExState, } from '../../../codewhisperer/views/lineAnnotationController' +import { AuthState } from '../../../auth/auth2' +import { createTestAuthUtil } from '../../testAuthUtil' describe('TryChatCodeLensProvider', () => { let instance: TryChatCodeLensProvider @@ -43,6 +45,7 @@ describe('TryChatCodeLensProvider', () => { }) beforeEach(async function () { + await createTestAuthUtil() isAmazonQVisibleEventEmitter = new vscode.EventEmitter() isAmazonQVisibleEvent = isAmazonQVisibleEventEmitter.event instance = new TryChatCodeLensProvider(isAmazonQVisibleEvent, () => codeLensPosition) @@ -58,7 +61,7 @@ describe('TryChatCodeLensProvider', () => { }) function stubConnection(state: AuthState) { - return sinon.stub(AuthUtil.instance, 'getChatAuthStateSync').returns({ amazonQ: state } as FeatureAuthState) + return sinon.stub(AuthUtil.instance, 'getAuthState').returns(state) } it('keeps returning a code lense until it hits the max times it should show', async function () { @@ -115,7 +118,9 @@ describe('TryChatCodeLensProvider', () => { stub.restore() } - const testStates = Object.values(AuthStates).filter((s) => s !== AuthStates.connected) + const testStates = Object.values(['connected', 'notConnected', 'expired'] as AuthState[]).filter( + (s) => s !== 'connected' + ) for (const state of testStates) { await testConnection(state) } diff --git a/packages/core/src/test/credentials/auth2.test.ts b/packages/core/src/test/credentials/auth2.test.ts new file mode 100644 index 00000000000..3f3df667d21 --- /dev/null +++ b/packages/core/src/test/credentials/auth2.test.ts @@ -0,0 +1,530 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as sinon from 'sinon' +import * as vscode from 'vscode' +import { LanguageClientAuth, SsoLogin } from '../../auth/auth2' +import { LanguageClient } from 'vscode-languageclient' +import { + GetSsoTokenResult, + SsoTokenSourceKind, + AuthorizationFlowKind, + ListProfilesResult, + UpdateCredentialsParams, + SsoTokenChangedParams, + bearerCredentialsUpdateRequestType, + bearerCredentialsDeleteNotificationType, + ssoTokenChangedRequestType, + SsoTokenChangedKind, + invalidateSsoTokenRequestType, + ProfileKind, + AwsErrorCodes, +} from '@aws/language-server-runtimes/protocol' +import * as ssoProvider from '../../auth/sso/ssoAccessTokenProvider' + +const profileName = 'test-profile' +const sessionName = 'test-session' +const region = 'us-east-1' +const startUrl = 'test-url' +const tokenId = 'test-token' + +describe('LanguageClientAuth', () => { + let client: sinon.SinonStubbedInstance + let auth: LanguageClientAuth + const encryptionKey = Buffer.from('test-key') + let useDeviceFlowStub: sinon.SinonStub + + beforeEach(() => { + client = sinon.createStubInstance(LanguageClient) + auth = new LanguageClientAuth(client as unknown as LanguageClient, 'testClient', encryptionKey) + useDeviceFlowStub = sinon.stub(ssoProvider, 'useDeviceFlow') + }) + + afterEach(() => { + sinon.restore() + }) + + describe('getSsoToken', () => { + async function testGetSsoToken(useDeviceFlow: boolean) { + const tokenSource = { + kind: SsoTokenSourceKind.IamIdentityCenter, + profileName, + } + useDeviceFlowStub.returns(useDeviceFlow ? true : false) + + await auth.getSsoToken(tokenSource, true) + + sinon.assert.calledOnce(client.sendRequest) + sinon.assert.calledWith( + client.sendRequest, + sinon.match.any, + sinon.match({ + clientName: 'testClient', + source: tokenSource, + options: { + loginOnInvalidToken: true, + authorizationFlow: useDeviceFlow + ? AuthorizationFlowKind.DeviceCode + : AuthorizationFlowKind.Pkce, + }, + }) + ) + } + + it('sends correct request parameters for pkce flow', async () => { + await testGetSsoToken(false) + }) + + it('sends correct request parameters for device code flow', async () => { + await testGetSsoToken(true) + }) + }) + + describe('updateProfile', () => { + it('sends correct profile update parameters', async () => { + await auth.updateProfile(profileName, startUrl, region, ['scope1']) + + sinon.assert.calledOnce(client.sendRequest) + const requestParams = client.sendRequest.firstCall.args[1] + sinon.assert.match(requestParams.profile, { + name: profileName, + }) + sinon.assert.match(requestParams.ssoSession.settings, { + sso_region: region, + }) + }) + }) + + describe('getProfile', () => { + const profile = { name: profileName, settings: { sso_session: sessionName } } + const ssoSession = { name: sessionName, settings: { sso_region: region, sso_start_url: startUrl } } + + it('returns the correct profile and sso session', async () => { + const mockListProfilesResult: ListProfilesResult = { + profiles: [ + { + ...profile, + kinds: [], + }, + ], + ssoSessions: [ssoSession], + } + client.sendRequest.resolves(mockListProfilesResult) + + const result = await auth.getProfile(profileName) + + sinon.assert.calledOnce(client.sendRequest) + sinon.assert.match(result, { + profile, + ssoSession, + }) + }) + + it('returns undefined for non-existent profile', async () => { + const mockListProfilesResult: ListProfilesResult = { + profiles: [], + ssoSessions: [], + } + client.sendRequest.resolves(mockListProfilesResult) + + const result = await auth.getProfile('non-existent-profile') + + sinon.assert.calledOnce(client.sendRequest) + sinon.assert.match(result, { profile: undefined, ssoSession: undefined }) + }) + }) + + describe('updateBearerToken', () => { + it('sends request', async () => { + const updateParams: UpdateCredentialsParams = { + data: 'token-data', + encrypted: true, + } + + await auth.updateBearerToken(updateParams) + + sinon.assert.calledOnce(client.sendRequest) + sinon.assert.calledWith(client.sendRequest, bearerCredentialsUpdateRequestType.method, updateParams) + }) + }) + + describe('deleteBearerToken', () => { + it('sends notification', async () => { + auth.deleteBearerToken() + + sinon.assert.calledOnce(client.sendNotification) + sinon.assert.calledWith(client.sendNotification, bearerCredentialsDeleteNotificationType.method) + }) + }) + + describe('invalidateSsoToken', () => { + it('sends request', async () => { + client.sendRequest.resolves({ success: true }) + const result = await auth.invalidateSsoToken(tokenId) + + sinon.assert.calledOnce(client.sendRequest) + sinon.assert.calledWith(client.sendRequest, invalidateSsoTokenRequestType.method, { ssoTokenId: tokenId }) + sinon.assert.match(result, { success: true }) + }) + }) + + describe('registerSsoTokenChangedHandler', () => { + it('registers the handler correctly', () => { + const handler = sinon.spy() + + auth.registerSsoTokenChangedHandler(handler) + + sinon.assert.calledOnce(client.onNotification) + sinon.assert.calledWith(client.onNotification, ssoTokenChangedRequestType.method, sinon.match.func) + + // Simulate a token changed notification + const tokenChangedParams: SsoTokenChangedParams = { + kind: SsoTokenChangedKind.Refreshed, + ssoTokenId: tokenId, + } + const registeredHandler = client.onNotification.firstCall.args[1] + registeredHandler(tokenChangedParams) + + sinon.assert.calledOnce(handler) + sinon.assert.calledWith(handler, tokenChangedParams) + }) + }) +}) + +describe('SsoLogin', () => { + let lspAuth: sinon.SinonStubbedInstance + let ssoLogin: SsoLogin + let eventEmitter: vscode.EventEmitter + let fireEventSpy: sinon.SinonSpy + + const loginOpts = { + startUrl, + region, + scopes: ['scope1'], + } + + const mockGetSsoTokenResponse: GetSsoTokenResult = { + ssoToken: { + id: tokenId, + accessToken: 'encrypted-token', + }, + updateCredentialsParams: { + data: '', + }, + } + + beforeEach(() => { + lspAuth = sinon.createStubInstance(LanguageClientAuth) + eventEmitter = new vscode.EventEmitter() + fireEventSpy = sinon.spy(eventEmitter, 'fire') + ssoLogin = new SsoLogin(profileName, lspAuth as any) + ;(ssoLogin as any).eventEmitter = eventEmitter + ;(ssoLogin as any).connectionState = 'notConnected' + }) + + afterEach(() => { + sinon.restore() + eventEmitter.dispose() + }) + + describe('login', () => { + it('updates profile and returns SSO token', async () => { + lspAuth.updateProfile.resolves() + lspAuth.getSsoToken.resolves(mockGetSsoTokenResponse) + + const response = await ssoLogin.login(loginOpts) + + sinon.assert.calledOnce(lspAuth.updateProfile) + sinon.assert.calledWith( + lspAuth.updateProfile, + profileName, + loginOpts.startUrl, + loginOpts.region, + loginOpts.scopes + ) + sinon.assert.calledOnce(lspAuth.getSsoToken) + sinon.assert.match(ssoLogin.getConnectionState(), 'connected') + sinon.assert.match(ssoLogin.data, { + startUrl: loginOpts.startUrl, + region: loginOpts.region, + }) + sinon.assert.match(response.ssoToken.id, tokenId) + sinon.assert.match(response.updateCredentialsParams, mockGetSsoTokenResponse.updateCredentialsParams) + }) + }) + + describe('reauthenticate', () => { + it('throws when not connected', async () => { + ;(ssoLogin as any).connectionState = 'notConnected' + try { + await ssoLogin.reauthenticate() + sinon.assert.fail('Should have thrown an error') + } catch (err) { + sinon.assert.match((err as Error).message, 'Cannot reauthenticate when not connected.') + } + }) + + it('returns new SSO token when connected', async () => { + ;(ssoLogin as any).connectionState = 'connected' + lspAuth.getSsoToken.resolves(mockGetSsoTokenResponse) + + const response = await ssoLogin.reauthenticate() + + sinon.assert.calledOnce(lspAuth.getSsoToken) + sinon.assert.match(ssoLogin.getConnectionState(), 'connected') + sinon.assert.match(response.ssoToken.id, tokenId) + sinon.assert.match(response.updateCredentialsParams, mockGetSsoTokenResponse.updateCredentialsParams) + }) + }) + + describe('logout', () => { + it('invalidates token and updates state', async () => { + await ssoLogin.logout() + + sinon.assert.match(ssoLogin.getConnectionState(), 'notConnected') + sinon.assert.match(ssoLogin.data, undefined) + }) + + it('emits state change event', async () => { + ;(ssoLogin as any).connectionState = 'connected' + ;(ssoLogin as any).ssoTokenId = tokenId + ;(ssoLogin as any)._data = { + startUrl: loginOpts.startUrl, + region: loginOpts.region, + } + ;(ssoLogin as any).eventEmitter = eventEmitter + + lspAuth.invalidateSsoToken.resolves({ success: true }) + + await ssoLogin.logout() + + sinon.assert.calledOnce(fireEventSpy) + sinon.assert.calledWith(fireEventSpy, { + id: profileName, + state: 'notConnected', + }) + }) + }) + + describe('restore', () => { + const mockProfile = { + profile: { + kinds: [ProfileKind.SsoTokenProfile], + name: profileName, + }, + ssoSession: { + name: sessionName, + settings: { + sso_region: region, + sso_start_url: startUrl, + }, + }, + } + + it('restores connection state from existing profile', async () => { + lspAuth.getProfile.resolves(mockProfile) + lspAuth.getSsoToken.resolves(mockGetSsoTokenResponse) + + await ssoLogin.restore() + + sinon.assert.calledOnce(lspAuth.getProfile) + sinon.assert.calledWith(lspAuth.getProfile, mockProfile.profile.name) + sinon.assert.calledOnce(lspAuth.getSsoToken) + sinon.assert.calledWith( + lspAuth.getSsoToken, + sinon.match({ + kind: SsoTokenSourceKind.IamIdentityCenter, + profileName: mockProfile.profile.name, + }), + false // login parameter + ) + + sinon.assert.match(ssoLogin.data, { + region: region, + startUrl: startUrl, + }) + sinon.assert.match(ssoLogin.getConnectionState(), 'connected') + sinon.assert.match((ssoLogin as any).ssoTokenId, tokenId) + }) + + it('does not connect for non-existent profile', async () => { + lspAuth.getProfile.resolves({ profile: undefined, ssoSession: undefined }) + + await ssoLogin.restore() + + sinon.assert.calledOnce(lspAuth.getProfile) + sinon.assert.calledOnce(lspAuth.getSsoToken) + sinon.assert.match(ssoLogin.data, undefined) + sinon.assert.match(ssoLogin.getConnectionState(), 'notConnected') + }) + + it('emits state change event on successful restore', async () => { + ;(ssoLogin as any).eventEmitter = eventEmitter + + lspAuth.getProfile.resolves(mockProfile) + lspAuth.getSsoToken.resolves(mockGetSsoTokenResponse) + + await ssoLogin.restore() + + sinon.assert.calledOnce(fireEventSpy) + sinon.assert.calledWith(fireEventSpy, { + id: profileName, + state: 'connected', + }) + }) + }) + + describe('cancelLogin', () => { + it('cancels and dispose token source', async () => { + await ssoLogin.login(loginOpts).catch(() => {}) + + ssoLogin.cancelLogin() + + const tokenSource = (ssoLogin as any).cancellationToken + sinon.assert.match(tokenSource, undefined) + }) + }) + + describe('_getSsoToken', () => { + beforeEach(() => { + ;(ssoLogin as any).connectionState = 'connected' + }) + + const testErrorHandling = async (errorCode: string, expectedState: string, shouldEmitEvent: boolean = true) => { + const error = new Error('Token error') + ;(error as any).data = { awsErrorCode: errorCode } + lspAuth.getSsoToken.rejects(error) + + try { + await (ssoLogin as any)._getSsoToken(false) + sinon.assert.fail('Should have thrown an error') + } catch (err) { + sinon.assert.match(err, error) + } + + sinon.assert.match(ssoLogin.getConnectionState(), expectedState) + + if (shouldEmitEvent) { + sinon.assert.calledWith(fireEventSpy, { + id: profileName, + state: expectedState, + }) + } + + sinon.assert.match((ssoLogin as any).cancellationToken, undefined) + } + + const notConnectedErrors = [ + AwsErrorCodes.E_CANCELLED, + AwsErrorCodes.E_SSO_SESSION_NOT_FOUND, + AwsErrorCodes.E_PROFILE_NOT_FOUND, + AwsErrorCodes.E_INVALID_SSO_TOKEN, + ] + + for (const errorCode of notConnectedErrors) { + it(`handles ${errorCode} error`, async () => { + await testErrorHandling(errorCode, 'notConnected') + }) + } + + it('handles token refresh error', async () => { + await testErrorHandling(AwsErrorCodes.E_CANNOT_REFRESH_SSO_TOKEN, 'expired') + }) + + it('handles unknown errors', async () => { + await testErrorHandling('UNKNOWN_ERROR', ssoLogin.getConnectionState(), false) + }) + + it('returns correct response and cleans up cancellation token', async () => { + lspAuth.getSsoToken.resolves(mockGetSsoTokenResponse) + + const response = await (ssoLogin as any)._getSsoToken(true) + + sinon.assert.calledWith( + lspAuth.getSsoToken, + sinon.match({ + kind: SsoTokenSourceKind.IamIdentityCenter, + profileName, + }), + true + ) + + sinon.assert.match(response, mockGetSsoTokenResponse) + sinon.assert.match((ssoLogin as any).cancellationToken, undefined) + }) + + it('updates state when token is retrieved successfully', async () => { + ;(ssoLogin as any).connectionState = 'notConnected' + lspAuth.getSsoToken.resolves(mockGetSsoTokenResponse) + + await (ssoLogin as any)._getSsoToken(true) + + sinon.assert.match(ssoLogin.getConnectionState(), 'connected') + sinon.assert.match((ssoLogin as any).ssoTokenId, tokenId) + sinon.assert.calledWith(fireEventSpy, { + id: profileName, + state: 'connected', + }) + }) + }) + + describe('onDidChangeConnectionState', () => { + it('should register handler for connection state changes', () => { + const handler = sinon.spy() + ssoLogin.onDidChangeConnectionState(handler) + + // Simulate state change + ;(ssoLogin as any).updateConnectionState('connected') + + sinon.assert.calledWith(handler, { + id: profileName, + state: 'connected', + }) + }) + }) + + describe('ssoTokenChangedHandler', () => { + beforeEach(() => { + ;(ssoLogin as any).ssoTokenId = tokenId + ;(ssoLogin as any).connectionState = 'connected' + }) + + it('updates state when token expires', () => { + ;(ssoLogin as any).ssoTokenChangedHandler({ + kind: 'Expired', + ssoTokenId: tokenId, + }) + + sinon.assert.match(ssoLogin.getConnectionState(), 'expired') + sinon.assert.calledOnce(fireEventSpy) + sinon.assert.calledWith(fireEventSpy, { + id: profileName, + state: 'expired', + }) + }) + + it('emits refresh event when token is refreshed', () => { + ;(ssoLogin as any).ssoTokenChangedHandler({ + kind: 'Refreshed', + ssoTokenId: tokenId, + }) + + sinon.assert.calledOnce(fireEventSpy) + sinon.assert.calledWith(fireEventSpy, { + id: profileName, + state: 'refreshed', + }) + }) + + it('does not emit event for different token ID', () => { + ;(ssoLogin as any).ssoTokenChangedHandler({ + kind: 'Refreshed', + ssoTokenId: 'different-token-id', + }) + + sinon.assert.notCalled(fireEventSpy) + }) + }) +}) diff --git a/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts b/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts index 2cb98193224..bd7f264f557 100644 --- a/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts +++ b/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts @@ -83,7 +83,7 @@ describe('SsoAccessTokenProvider', function () { tempDir = await makeTemporaryTokenCacheFolder() cache = getCache(tempDir) reAuthState = new TestReAuthState() - sut = SsoAccessTokenProvider.create({ region, startUrl }, cache, oidcClient, reAuthState, () => true) + sut = SsoAccessTokenProvider.create({ region, startUrl }, cache, oidcClient, reAuthState) }) afterEach(async function () { @@ -271,13 +271,7 @@ describe('SsoAccessTokenProvider', function () { await sut.createToken() // Mimic when we sign out then in again with the same region+startUrl. The ID is the only thing different. - sut = SsoAccessTokenProvider.create( - { region, startUrl, identifier: 'bbb' }, - cache, - oidcClient, - reAuthState, - () => true - ) + sut = SsoAccessTokenProvider.create({ region, startUrl, identifier: 'bbb' }, cache, oidcClient, reAuthState) await sut.createToken() assertTelemetry('aws_loginWithBrowser', [ diff --git a/packages/core/src/test/credentials/utils.test.ts b/packages/core/src/test/credentials/utils.test.ts index ff4dc1046bc..2526118d5e1 100644 --- a/packages/core/src/test/credentials/utils.test.ts +++ b/packages/core/src/test/credentials/utils.test.ts @@ -67,7 +67,6 @@ type SsoTestCase = { kind: SsoKind; connections: Connection[]; expected: boolean type BuilderIdTestCase = { kind: BuilderIdKind; connections: Connection[]; expected: boolean } describe('connection exists funcs', function () { - const cwIdcConnection: SsoConnection = { ...ssoConnection, scopes: amazonQScopes, label: 'codeWhispererSso' } const cwBuilderIdConnection: SsoConnection = { ...builderIdConnection, scopes: amazonQScopes, @@ -81,7 +80,6 @@ describe('connection exists funcs', function () { const ssoConnections: Connection[] = [ ssoConnection, builderIdConnection, - cwIdcConnection, cwBuilderIdConnection, ccBuilderIdConnection, ] @@ -96,15 +94,13 @@ describe('connection exists funcs', function () { ].map((c) => { return { ...c, kind: 'any' } }) - const cwIdcCases: SsoTestCase[] = [ - { connections: [cwIdcConnection], expected: true }, - { connections: allConnections, expected: true }, + const ccIdcCases: SsoTestCase[] = [ { connections: [], expected: false }, - { connections: allConnections.filter((c) => c !== cwIdcConnection), expected: false }, + { connections: allConnections, expected: false }, ].map((c) => { - return { ...c, kind: 'codewhisperer' } + return { ...c, kind: 'codecatalyst' } }) - const allCases = [...anyCases, ...cwIdcCases] + const allCases = [...anyCases, ...ccIdcCases] for (const args of allCases) { it(`ssoExists() returns '${args.expected}' when kind '${args.kind}' given [${args.connections @@ -116,15 +112,6 @@ describe('connection exists funcs', function () { }) describe('builderIdExists()', function () { - const cwBuilderIdCases: BuilderIdTestCase[] = [ - { connections: [cwBuilderIdConnection], expected: true }, - { connections: allConnections, expected: true }, - { connections: [], expected: false }, - { connections: allConnections.filter((c) => c !== cwBuilderIdConnection), expected: false }, - ].map((c) => { - return { ...c, kind: 'codewhisperer' } - }) - const ccBuilderIdCases: BuilderIdTestCase[] = [ { connections: [ccBuilderIdConnection], expected: true }, { connections: allConnections, expected: true }, @@ -134,9 +121,7 @@ describe('connection exists funcs', function () { return { ...c, kind: 'codecatalyst' } }) - const allCases = [...cwBuilderIdCases, ...ccBuilderIdCases] - - for (const args of allCases) { + for (const args of ccBuilderIdCases) { it(`builderIdExists() returns '${args.expected}' when kind '${args.kind}' given [${args.connections .map((c) => c.label) .join(', ')}]`, async function () { diff --git a/packages/core/src/test/index.ts b/packages/core/src/test/index.ts index 9a01973e26d..0a4203f2fe5 100644 --- a/packages/core/src/test/index.ts +++ b/packages/core/src/test/index.ts @@ -25,3 +25,4 @@ export * from './testUtil' export * from './amazonq/utils' export * from './fake/mockFeatureConfigData' export * from './shared/ui/testUtils' +export * from './testAuthUtil' diff --git a/packages/core/src/test/login/webview/vue/backend_amazonq.test.ts b/packages/core/src/test/login/webview/vue/backend_amazonq.test.ts deleted file mode 100644 index c22b1a77fe1..00000000000 --- a/packages/core/src/test/login/webview/vue/backend_amazonq.test.ts +++ /dev/null @@ -1,149 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import { SinonSandbox, createSandbox } from 'sinon' -import { assertTelemetry } from '../../../testUtil' -import assert from 'assert' -import { - createBuilderIdProfile, - createSsoProfile, - createTestAuth, - mockRegistration, -} from '../../../credentials/testUtil' -import { Auth } from '../../../../auth' -import { AmazonQLoginWebview } from '../../../../login/webview/vue/amazonq/backend_amazonq' -import { isBuilderIdConnection, isIdcSsoConnection } from '../../../../auth/connection' -import { amazonQScopes, AuthUtil } from '../../../../codewhisperer/util/authUtil' -import { getOpenExternalStub } from '../../../globalSetup.test' -import globals from '../../../../shared/extensionGlobals' - -// TODO: remove auth page and tests -describe('Amazon Q Login', function () { - const region = 'fakeRegion' - const startUrl = 'fakeUrl' - - let sandbox: SinonSandbox - let auth: ReturnType - let authUtil: AuthUtil - let backend: AmazonQLoginWebview - - beforeEach(function () { - sandbox = createSandbox() - auth = createTestAuth(globals.globalState) - authUtil = new AuthUtil(auth) - sandbox.stub(Auth, 'instance').value(auth) - sandbox.stub(AuthUtil, 'instance').value(authUtil) - getOpenExternalStub().resolves(true) - - backend = new AmazonQLoginWebview() - }) - - afterEach(function () { - sandbox.restore() - }) - - it('signs into builder ID and emits telemetry', async function () { - await backend.startBuilderIdSetup() - - assert.ok(isBuilderIdConnection(auth.activeConnection)) - assert.deepStrictEqual(auth.activeConnection.scopes, amazonQScopes) - assert.deepStrictEqual(auth.activeConnection.state, 'valid') - - assertTelemetry('auth_addConnection', { - result: 'Succeeded', - credentialSourceId: 'awsId', - authEnabledFeatures: 'codewhisperer', - isReAuth: false, - ssoRegistrationExpiresAt: mockRegistration.expiresAt.toISOString(), - ssoRegistrationClientId: mockRegistration.clientId, - }) - }) - - it('signs into IdC and emits telemetry', async function () { - await backend.startEnterpriseSetup(startUrl, region) - - assert.ok(isIdcSsoConnection(auth.activeConnection)) - assert.deepStrictEqual(auth.activeConnection.scopes, amazonQScopes) - assert.deepStrictEqual(auth.activeConnection.state, 'valid') - assert.deepStrictEqual(auth.activeConnection.startUrl, startUrl) - assert.deepStrictEqual(auth.activeConnection.ssoRegion, region) - - assertTelemetry('auth_addConnection', { - result: 'Succeeded', - credentialSourceId: 'iamIdentityCenter', - authEnabledFeatures: 'codewhisperer', - credentialStartUrl: startUrl, - awsRegion: region, - isReAuth: false, - ssoRegistrationExpiresAt: mockRegistration.expiresAt.toISOString(), - ssoRegistrationClientId: mockRegistration.clientId, - }) - }) - - it('reauths builder ID and emits telemetry', async function () { - const conn = await auth.createInvalidSsoConnection(createBuilderIdProfile({ scopes: amazonQScopes })) - await auth.useConnection(conn) - - // method under test - await backend.reauthenticateConnection() - - assert.deepStrictEqual(auth.activeConnection?.state, 'valid') - - assertTelemetry('auth_addConnection', { - result: 'Succeeded', - credentialSourceId: 'awsId', - authEnabledFeatures: 'codewhisperer', - isReAuth: true, - ssoRegistrationExpiresAt: mockRegistration.expiresAt.toISOString(), - ssoRegistrationClientId: mockRegistration.clientId, - }) - }) - - it('reauths IdC and emits telemetry', async function () { - const conn = await auth.createInvalidSsoConnection( - createSsoProfile({ scopes: amazonQScopes, startUrl, ssoRegion: region }) - ) - await auth.useConnection(conn) - - // method under test - await backend.reauthenticateConnection() - - assert.deepStrictEqual(auth.activeConnection?.state, 'valid') - - assertTelemetry('auth_addConnection', { - result: 'Succeeded', - credentialSourceId: 'iamIdentityCenter', - authEnabledFeatures: 'codewhisperer', - credentialStartUrl: startUrl, - awsRegion: region, - isReAuth: true, - ssoRegistrationExpiresAt: mockRegistration.expiresAt.toISOString(), - ssoRegistrationClientId: mockRegistration.clientId, - }) - }) - - it('signs out of reauth and emits telemetry', async function () { - const conn = await auth.createInvalidSsoConnection( - createSsoProfile({ scopes: amazonQScopes, startUrl, ssoRegion: region }) - ) - await auth.useConnection(conn) - - // method under test - await backend.signout() - - assert.equal(auth.activeConnection, undefined) - - assertTelemetry('auth_addConnection', { - result: 'Cancelled', - credentialSourceId: 'iamIdentityCenter', - authEnabledFeatures: 'codewhisperer', - credentialStartUrl: startUrl, - awsRegion: region, - isReAuth: true, - ssoRegistrationExpiresAt: mockRegistration.expiresAt.toISOString(), - ssoRegistrationClientId: mockRegistration.clientId, - }) - }) -}) diff --git a/packages/core/src/test/shared/featureConfig.test.ts b/packages/core/src/test/shared/featureConfig.test.ts index 9c9f20cb5fb..e94358361b0 100644 --- a/packages/core/src/test/shared/featureConfig.test.ts +++ b/packages/core/src/test/shared/featureConfig.test.ts @@ -10,9 +10,12 @@ import { Features, FeatureConfigProvider, featureDefinitions, FeatureName } from import { ListFeatureEvaluationsResponse } from '../../codewhisperer' import { createSpyClient } from '../codewhisperer/testUtil' import { mockFeatureConfigsData } from '../fake/mockFeatureConfigData' +import { createTestAuthUtil } from '../testAuthUtil' describe('FeatureConfigProvider', () => { beforeEach(async () => { + await createTestAuthUtil() + const clientSpy = await createSpyClient() sinon.stub(clientSpy, 'listFeatureEvaluations').returns({ promise: () => diff --git a/packages/core/src/test/shared/utilities/functionUtils.test.ts b/packages/core/src/test/shared/utilities/functionUtils.test.ts index 7880d11ff63..b675fe74feb 100644 --- a/packages/core/src/test/shared/utilities/functionUtils.test.ts +++ b/packages/core/src/test/shared/utilities/functionUtils.test.ts @@ -152,6 +152,33 @@ describe('debounce', function () { assert.strictEqual(counter, 2) }) + describe('useLastCall option', function () { + let args: number[] + let clock: ReturnType + let addToArgs: (i: number) => void + + before(function () { + args = [] + clock = installFakeClock() + addToArgs = (n: number) => args.push(n) + }) + + afterEach(function () { + clock.uninstall() + args.length = 0 + }) + + it('only calls with the last args', async function () { + const debounced = debounce(addToArgs, 10, true) + const p1 = debounced(1) + const p2 = debounced(2) + const p3 = debounced(3) + await clock.tickAsync(100) + await Promise.all([p1, p2, p3]) + assert.deepStrictEqual(args, [3]) + }) + }) + describe('window rolling', function () { let clock: ReturnType const calls: ReturnType[] = [] diff --git a/packages/core/src/test/testAuthUtil.ts b/packages/core/src/test/testAuthUtil.ts new file mode 100644 index 00000000000..595f8bf45ef --- /dev/null +++ b/packages/core/src/test/testAuthUtil.ts @@ -0,0 +1,47 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as sinon from 'sinon' +import * as jose from 'jose' +import * as crypto from 'crypto' +import { LanguageClientAuth } from '../auth/auth2' +import { AuthUtil } from '../codewhisperer/util/authUtil' + +export async function createTestAuthUtil() { + const encryptionKey = crypto.randomBytes(32) + + const jwe = await new jose.CompactEncrypt(new TextEncoder().encode(JSON.stringify({ your: 'mock data' }))) + .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' }) + .encrypt(encryptionKey) + + const fakeToken = { + ssoToken: { + id: 'fake-id', + accessToken: jwe, + }, + updateCredentialsParams: { + data: 'fake-data', + }, + } + + const mockLspAuth: Partial = { + registerSsoTokenChangedHandler: sinon.stub().resolves(), + updateProfile: sinon.stub().resolves(), + getSsoToken: sinon.stub().resolves(fakeToken), + getProfile: sinon.stub().resolves({ + sso_registration_scopes: ['codewhisperer'], + }), + deleteBearerToken: sinon.stub().resolves(), + updateBearerToken: sinon.stub().resolves(), + invalidateSsoToken: sinon.stub().resolves(), + registerCacheWatcher: sinon.stub().resolves(), + encryptionKey, + } + + // Since AuthUtil is a singleton, we want to remove an existing instance before setting up a new one + AuthUtil.destroy() + + AuthUtil.create(mockLspAuth as LanguageClientAuth) +} diff --git a/packages/core/src/testE2E/codewhisperer/referenceTracker.test.ts b/packages/core/src/testE2E/codewhisperer/referenceTracker.test.ts deleted file mode 100644 index 0038795ad89..00000000000 --- a/packages/core/src/testE2E/codewhisperer/referenceTracker.test.ts +++ /dev/null @@ -1,125 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as codewhispererClient from '../../codewhisperer/client/codewhisperer' -import { ConfigurationEntry } from '../../codewhisperer/models/model' -import { setValidConnection, skipTestIfNoValidConn } from '../util/connection' -import { RecommendationHandler } from '../../codewhisperer/service/recommendationHandler' -import { createMockTextEditor, resetCodeWhispererGlobalVariables } from '../../test/codewhisperer/testUtil' -import { invokeRecommendation } from '../../codewhisperer/commands/invokeRecommendation' -import { session } from '../../codewhisperer/util/codeWhispererSession' - -/* -New model deployment may impact references returned. - -These tests: - 1) are not required for github approval flow - 2) will be auto-skipped until fix for manual runs is posted. -*/ - -const leftContext = `InAuto.GetContent( - InAuto.servers.auto, "vendors.json", - function (data) { - let block = ''; - for(let i = 0; i < data.length; i++) { - block += '' + cars[i].title + ''; - } - $('#cars').html(block); - });` - -describe('CodeWhisperer service invocation', async function () { - let validConnection: boolean - const client = new codewhispererClient.DefaultCodeWhispererClient() - const configWithRefs: ConfigurationEntry = { - isShowMethodsEnabled: true, - isManualTriggerEnabled: true, - isAutomatedTriggerEnabled: true, - isSuggestionsWithCodeReferencesEnabled: true, - } - const configWithNoRefs: ConfigurationEntry = { - isShowMethodsEnabled: true, - isManualTriggerEnabled: true, - isAutomatedTriggerEnabled: true, - isSuggestionsWithCodeReferencesEnabled: false, - } - - before(async function () { - validConnection = await setValidConnection() - }) - - beforeEach(function () { - void resetCodeWhispererGlobalVariables() - RecommendationHandler.instance.clearRecommendations() - // TODO: remove this line (this.skip()) when these tests no longer auto-skipped - this.skip() - // valid connection required to run tests - skipTestIfNoValidConn(validConnection, this) - }) - - it('trigger known to return recs with references returns rec with reference', async function () { - // check that handler is empty before invocation - const requestIdBefore = RecommendationHandler.instance.requestId - const sessionIdBefore = session.sessionId - const validRecsBefore = RecommendationHandler.instance.isValidResponse() - - assert.ok(requestIdBefore.length === 0) - assert.ok(sessionIdBefore.length === 0) - assert.ok(!validRecsBefore) - - const doc = leftContext + rightContext - const filename = 'test.js' - const language = 'javascript' - const line = 5 - const character = 39 - const mockEditor = createMockTextEditor(doc, filename, language, line, character) - - await invokeRecommendation(mockEditor, client, configWithRefs) - - const requestId = RecommendationHandler.instance.requestId - const sessionId = session.sessionId - const validRecs = RecommendationHandler.instance.isValidResponse() - const references = session.recommendations[0].references - - assert.ok(requestId.length > 0) - assert.ok(sessionId.length > 0) - assert.ok(validRecs) - assert.ok(references !== undefined) - // TODO: uncomment this assert when this test is no longer auto-skipped - // assert.ok(references.length > 0) - }) - - // This test will fail if user is logged in with IAM identity center - it('trigger known to return rec with references does not return rec with references when reference tracker setting is off', async function () { - // check that handler is empty before invocation - const requestIdBefore = RecommendationHandler.instance.requestId - const sessionIdBefore = session.sessionId - const validRecsBefore = RecommendationHandler.instance.isValidResponse() - - assert.ok(requestIdBefore.length === 0) - assert.ok(sessionIdBefore.length === 0) - assert.ok(!validRecsBefore) - - const doc = leftContext + rightContext - const filename = 'test.js' - const language = 'javascript' - const line = 5 - const character = 39 - const mockEditor = createMockTextEditor(doc, filename, language, line, character) - - await invokeRecommendation(mockEditor, client, configWithNoRefs) - - const requestId = RecommendationHandler.instance.requestId - const sessionId = session.sessionId - const validRecs = RecommendationHandler.instance.isValidResponse() - - assert.ok(requestId.length > 0) - assert.ok(sessionId.length > 0) - // no recs returned because example request returns 1 rec with reference, so no recs returned when references off - assert.ok(!validRecs) - }) -}) diff --git a/packages/core/src/testE2E/codewhisperer/securityScan.test.ts b/packages/core/src/testE2E/codewhisperer/securityScan.test.ts index 730b9628290..0dc5c5f105d 100644 --- a/packages/core/src/testE2E/codewhisperer/securityScan.test.ts +++ b/packages/core/src/testE2E/codewhisperer/securityScan.test.ts @@ -9,7 +9,7 @@ import * as codewhispererClient from '../../codewhisperer/client/codewhisperer' import * as CodeWhispererConstants from '../../codewhisperer/models/constants' import * as path from 'path' import * as testutil from '../../test/testUtil' -import { setValidConnection, skipTestIfNoValidConn } from '../util/connection' +import { skipTestIfNoValidConn } from '../util/connection' import { resetCodeWhispererGlobalVariables } from '../../test/codewhisperer/testUtil' import { getTestWorkspaceFolder } from '../../testInteg/integrationTestsUtilities' import { closeAllEditors } from '../../test/testUtil' @@ -23,6 +23,7 @@ import { makeTemporaryToolkitFolder } from '../../shared/filesystemUtilities' import fs from '../../shared/fs/fs' import { ZipUtil } from '../../codewhisperer/util/zipUtil' import { randomUUID } from '../../shared/crypto' +import { AuthUtil } from '../../codewhisperer' const filePromptWithSecurityIssues = `from flask import app @@ -53,7 +54,7 @@ describe('CodeWhisperer security scan', async function () { const workspaceFolder = getTestWorkspaceFolder() before(async function () { - validConnection = await setValidConnection() + validConnection = AuthUtil.instance.isConnected() }) beforeEach(function () { diff --git a/packages/core/src/testE2E/codewhisperer/serviceInvocations.test.ts b/packages/core/src/testE2E/codewhisperer/serviceInvocations.test.ts deleted file mode 100644 index d4265d13982..00000000000 --- a/packages/core/src/testE2E/codewhisperer/serviceInvocations.test.ts +++ /dev/null @@ -1,124 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import assert from 'assert' -import * as vscode from 'vscode' -import * as path from 'path' -import { setValidConnection, skipTestIfNoValidConn } from '../util/connection' -import { ConfigurationEntry } from '../../codewhisperer/models/model' -import * as codewhispererClient from '../../codewhisperer/client/codewhisperer' -import { RecommendationHandler } from '../../codewhisperer/service/recommendationHandler' -import { - createMockTextEditor, - createTextDocumentChangeEvent, - resetCodeWhispererGlobalVariables, -} from '../../test/codewhisperer/testUtil' -import { KeyStrokeHandler } from '../../codewhisperer/service/keyStrokeHandler' -import { sleep } from '../../shared/utilities/timeoutUtils' -import { invokeRecommendation } from '../../codewhisperer/commands/invokeRecommendation' -import { getTestWorkspaceFolder } from '../../testInteg/integrationTestsUtilities' -import { session } from '../../codewhisperer/util/codeWhispererSession' - -describe('CodeWhisperer service invocation', async function () { - let validConnection: boolean - const client = new codewhispererClient.DefaultCodeWhispererClient() - const config: ConfigurationEntry = { - isShowMethodsEnabled: true, - isManualTriggerEnabled: true, - isAutomatedTriggerEnabled: true, - isSuggestionsWithCodeReferencesEnabled: true, - } - - before(async function () { - validConnection = await setValidConnection() - }) - - beforeEach(function () { - void resetCodeWhispererGlobalVariables() - RecommendationHandler.instance.clearRecommendations() - // valid connection required to run tests - skipTestIfNoValidConn(validConnection, this) - }) - - it('manual trigger returns valid recommendation response', async function () { - // check that handler is empty before invocation - const requestIdBefore = RecommendationHandler.instance.requestId - const sessionIdBefore = session.sessionId - const validRecsBefore = RecommendationHandler.instance.isValidResponse() - - assert.ok(requestIdBefore.length === 0) - assert.ok(sessionIdBefore.length === 0) - assert.ok(!validRecsBefore) - - const mockEditor = createMockTextEditor() - await invokeRecommendation(mockEditor, client, config) - - const requestId = RecommendationHandler.instance.requestId - const sessionId = session.sessionId - const validRecs = RecommendationHandler.instance.isValidResponse() - - assert.ok(requestId.length > 0) - assert.ok(sessionId.length > 0) - assert.ok(validRecs) - }) - - it('auto trigger returns valid recommendation response', async function () { - // check that handler is empty before invocation - const requestIdBefore = RecommendationHandler.instance.requestId - const sessionIdBefore = session.sessionId - const validRecsBefore = RecommendationHandler.instance.isValidResponse() - - assert.ok(requestIdBefore.length === 0) - assert.ok(sessionIdBefore.length === 0) - assert.ok(!validRecsBefore) - - const mockEditor = createMockTextEditor() - - const mockEvent: vscode.TextDocumentChangeEvent = createTextDocumentChangeEvent( - mockEditor.document, - new vscode.Range(new vscode.Position(0, 0), new vscode.Position(0, 1)), - '\n' - ) - - await KeyStrokeHandler.instance.processKeyStroke(mockEvent, mockEditor, client, config) - // wait for 5 seconds to allow time for response to be generated - await sleep(5000) - - const requestId = RecommendationHandler.instance.requestId - const sessionId = session.sessionId - const validRecs = RecommendationHandler.instance.isValidResponse() - - assert.ok(requestId.length > 0) - assert.ok(sessionId.length > 0) - assert.ok(validRecs) - }) - - it('invocation in unsupported language does not generate a request', async function () { - const workspaceFolder = getTestWorkspaceFolder() - const appRoot = path.join(workspaceFolder, 'go1-plain-sam-app') - const appCodePath = path.join(appRoot, 'hello-world', 'main.go') - - // check that handler is empty before invocation - const requestIdBefore = RecommendationHandler.instance.requestId - const sessionIdBefore = session.sessionId - const validRecsBefore = RecommendationHandler.instance.isValidResponse() - - assert.ok(requestIdBefore.length === 0) - assert.ok(sessionIdBefore.length === 0) - assert.ok(!validRecsBefore) - - const doc = await vscode.workspace.openTextDocument(vscode.Uri.file(appCodePath)) - const editor = await vscode.window.showTextDocument(doc) - await invokeRecommendation(editor, client, config) - - const requestId = RecommendationHandler.instance.requestId - const sessionId = session.sessionId - const validRecs = RecommendationHandler.instance.isValidResponse() - - assert.ok(requestId.length === 0) - assert.ok(sessionId.length === 0) - assert.ok(!validRecs) - }) -}) diff --git a/packages/core/src/testE2E/util/connection.ts b/packages/core/src/testE2E/util/connection.ts index bf158426a2a..da2f733a458 100644 --- a/packages/core/src/testE2E/util/connection.ts +++ b/packages/core/src/testE2E/util/connection.ts @@ -3,9 +3,6 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { isValidAmazonQConnection } from '../../codewhisperer/util/authUtil' -import { Auth } from '../../auth/auth' - /* In order to run codewhisperer or gumby integration tests user must: @@ -16,24 +13,6 @@ Test cases will skip if the above criteria are not met. If user has an expired connection they must reauthenticate prior to running tests. */ -async function getValidConnection() { - return (await Auth.instance.listConnections()).find(isValidAmazonQConnection) -} - -export async function setValidConnection() { - const conn = await getValidConnection() - let validConnection: boolean - - if (conn !== undefined && Auth.instance.getConnectionState(conn) === 'valid') { - validConnection = true - await Auth.instance.useConnection(conn) - } else { - validConnection = false - console.log(`No valid auth connection, will skip Amazon Q integration test cases`) - } - return validConnection -} - export function skipTestIfNoValidConn(validConnection: boolean, ctx: Mocha.Context) { if (!validConnection && ctx.currentTest) { ctx.currentTest.title += ` (skipped - no valid connection)` diff --git a/packages/core/src/testInteg/amazonQTransform/transformByQ.test.ts b/packages/core/src/testInteg/amazonQTransform/transformByQ.test.ts index b191997a236..e2b423f0591 100644 --- a/packages/core/src/testInteg/amazonQTransform/transformByQ.test.ts +++ b/packages/core/src/testInteg/amazonQTransform/transformByQ.test.ts @@ -13,18 +13,16 @@ import request from '../../shared/request' import { transformByQState, ZipManifest } from '../../codewhisperer/models/model' import globals from '../../shared/extensionGlobals' import { fs } from '../../shared' -import { setValidConnection } from '../../testE2E/util/connection' +import { AuthUtil } from '../../codewhisperer/util/authUtil' describe('transformByQ', async function () { let tempDir = '' let tempFileName = '' let tempFilePath = '' let zippedCodePath = '' - let validConnection: boolean before(async function () { - validConnection = await setValidConnection() - if (!validConnection) { + if (!AuthUtil.instance.isConnected()) { this.skip() } tempDir = path.join(os.tmpdir(), 'gumby-test') diff --git a/packages/core/src/testInteg/perf/buildIndex.test.ts b/packages/core/src/testInteg/perf/buildIndex.test.ts deleted file mode 100644 index d60de3bdc3a..00000000000 --- a/packages/core/src/testInteg/perf/buildIndex.test.ts +++ /dev/null @@ -1,79 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import { performanceTest } from '../../shared/performance/performance' -import * as sinon from 'sinon' -import * as vscode from 'vscode' -import assert from 'assert' -import { LspClient, LspController } from '../../amazonq' -import { LanguageClient, ServerOptions } from 'vscode-languageclient' -import { createTestWorkspace } from '../../test/testUtil' -import { BuildIndexRequestType, GetUsageRequestType } from '../../amazonq/lsp/types' -import { fs, getRandomString } from '../../shared' -import { FileSystem } from '../../shared/fs/fs' -import { getFsCallsUpperBound } from './utilities' - -interface SetupResult { - clientReqStub: sinon.SinonStub - fsSpy: sinon.SinonSpiedInstance - findFilesSpy: sinon.SinonSpy -} - -async function verifyResult(setup: SetupResult) { - // A correct run makes 2 requests, but don't want to make it exact to avoid over-sensitivity to implementation. If we make 10+ something is likely wrong. - assert.ok(setup.clientReqStub.callCount >= 2 && setup.clientReqStub.callCount <= 10) - assert.ok(setup.clientReqStub.calledWith(BuildIndexRequestType)) - assert.ok(setup.clientReqStub.calledWith(GetUsageRequestType)) - - assert.strictEqual(getFsCallsUpperBound(setup.fsSpy), 0, 'should not make any fs calls') - assert.ok(setup.findFilesSpy.callCount <= 2, 'findFiles should not be called more than twice') -} - -async function setupWithWorkspace(numFiles: number, options: { fileContent: string }): Promise { - // Force VSCode to find my test workspace only to keep test contained and controlled. - const testWorksapce = await createTestWorkspace(numFiles, options) - sinon.stub(vscode.workspace, 'workspaceFolders').value([testWorksapce]) - - // Avoid sending real request to lsp. - const clientReqStub = sinon.stub(LanguageClient.prototype, 'sendRequest').resolves(true) - const fsSpy = sinon.spy(fs) - const findFilesSpy = sinon.spy(vscode.workspace, 'findFiles') - LspClient.instance.client = new LanguageClient('amazonq', 'test-client', {} as ServerOptions, {}) - return { clientReqStub, fsSpy, findFilesSpy } -} - -describe('buildIndex', function () { - describe('performanceTests', function () { - afterEach(function () { - sinon.restore() - }) - performanceTest({}, 'indexing many small files', function () { - return { - setup: async () => setupWithWorkspace(250, { fileContent: '0123456789' }), - execute: async () => { - await LspController.instance.buildIndex({ - startUrl: '', - maxIndexSize: 30, - isVectorIndexEnabled: true, - }) - }, - verify: verifyResult, - } - }) - performanceTest({}, 'indexing few large files', function () { - return { - setup: async () => setupWithWorkspace(10, { fileContent: getRandomString(1000) }), - execute: async () => { - await LspController.instance.buildIndex({ - startUrl: '', - maxIndexSize: 30, - isVectorIndexEnabled: true, - }) - }, - verify: verifyResult, - } - }) - }) -}) diff --git a/packages/toolkit/package.json b/packages/toolkit/package.json index 309f1e92da2..e99c8325388 100644 --- a/packages/toolkit/package.json +++ b/packages/toolkit/package.json @@ -255,6 +255,10 @@ "type": "boolean", "default": false }, + "amazonqLSPInlineChat": { + "type": "boolean", + "default": false + }, "amazonqChatLSP": { "type": "boolean", "default": true