Add support for device_code flow by selecting a checkbox

Joseph de Clerck · Joseph de Clerck · commit b2662ca59b3a · 2025-05-23T10:57:32.000-04:00
diff --git a/packages/core/src/auth/sso/model.ts b/packages/core/src/auth/sso/model.ts
@@ -14,6 +14,7 @@ import { getLogger } from '../../shared/logger/logger'
 import { CancellationError } from '../../shared/utilities/timeoutUtils'
 import { ssoAuthHelpUrl } from '../../shared/constants'
 import { openUrl } from '../../shared/utilities/vsCodeUtils'
+import { copyToClipboard } from '../../shared/utilities/messages'
 import { ToolkitError } from '../../shared/errors'
 import { builderIdStartUrl } from './constants'
 
@@ -100,6 +101,7 @@ export function truncateStartUrl(startUrl: string) {
 type Authorization = { readonly verificationUri: string; readonly userCode: string }
 
 export const proceedToBrowser = localize('AWS.auth.loginWithBrowser.proceedToBrowser', 'Proceed To Browser')
+export const copyUrl = localize('AWS.auth.loginWithBrowser.copyLoginUrl', 'Copy authentication URL')
 
 export async function openSsoPortalLink(startUrl: string, authorization: Authorization): Promise<boolean> {
     /**
@@ -122,13 +124,18 @@ export async function openSsoPortalLink(startUrl: string, authorization: Authori
             authorization.userCode
         )
 
+        const options = [proceedToBrowser, copyUrl]
+
         while (true) {
             // TODO: add the 'Help' item back once we have a suitable URL
             // const resp = await vscode.window.showInformationMessage(title, options, copyCode, localizedText.help)
-            const resp = await vscode.window.showInformationMessage(title, { modal: true, detail }, proceedToBrowser)
+            const resp = await vscode.window.showInformationMessage(title, { modal: true, detail }, ...options)
             switch (resp) {
                 case proceedToBrowser:
                     return openSsoUrl(makeConfirmCodeUrl(authorization))
+                case copyUrl:
+                    await copyToClipboard(makeConfirmCodeUrl(authorization).toString(true))
+                    return true
                 case localizedText.help:
                     await tryOpenHelpUrl(ssoAuthHelpUrl)
                     continue
diff --git a/packages/core/src/auth/sso/ssoAccessTokenProvider.ts b/packages/core/src/auth/sso/ssoAccessTokenProvider.ts
@@ -295,9 +295,15 @@ export abstract class SsoAccessTokenProvider {
              * 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)
+             * 3. User has explicitly requested device code flow via the login page toggle
              *
-             * Since we are unable to serve the final authorization page
+             * Since we are unable to serve the final authorization page in cases 1 and 2,
+             * or the user may not have federation access in case 3
              */
+            // Check for user preference first, then fall back to the current logic
+            if (globals.globalState.get<boolean>('aws.forceDeviceCodeFlow')) {
+                return true
+            }
             return getExtRuntimeContext().extensionHost === 'remote'
         }
     ) {
@@ -458,7 +464,8 @@ export class DeviceFlowAuthorization extends SsoAccessTokenProvider {
             clientId: registration.clientId,
             clientSecret: registration.clientSecret,
         })
-
+        // Reset forceDeviceCodeFlow if toggled
+        await globals.globalState.update('aws.forceDeviceCodeFlow', undefined)
         const openBrowserAndWaitUntilComplete = async () => {
             if (!(await openSsoPortalLink(this.profile.startUrl, authorization))) {
                 throw new CancellationError('user')
diff --git a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts
@@ -78,12 +78,24 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
         })
     }
 
-    async startEnterpriseSetup(startUrl: string, region: string): Promise<AuthError | undefined> {
-        getLogger().debug(`called startEnterpriseSetup() with startUrl: '${startUrl}', region: '${region}'`)
+    async startEnterpriseSetup(
+        startUrl: string,
+        region: string,
+        app: string,
+        useDeviceCodeFlow?: boolean
+    ): Promise<AuthError | undefined> {
+        getLogger().debug(
+            `called startEnterpriseSetup() with startUrl: '${startUrl}', region: '${region}', useDeviceCodeFlow: ${useDeviceCodeFlow}`
+        )
         await globals.globalState.update('recentSso', {
             startUrl: startUrl,
             region: region,
         })
+        // Store the device code flow preference if it's set
+        if (useDeviceCodeFlow) {
+            await globals.globalState.update('aws.forceDeviceCodeFlow', true)
+        }
+
         return await this.ssoSetup('startCodeWhispererEnterpriseSetup', async () => {
             this.storeMetricMetadata({
                 credentialStartUrl: startUrl,
diff --git a/packages/core/src/login/webview/vue/backend.ts b/packages/core/src/login/webview/vue/backend.ts
@@ -165,7 +165,12 @@ export abstract class CommonAuthWebview extends VueWebview {
 
     abstract startBuilderIdSetup(app: string): Promise<AuthError | undefined>
 
-    abstract startEnterpriseSetup(startUrl: string, region: string, app: string): Promise<AuthError | undefined>
+    abstract startEnterpriseSetup(
+        startUrl: string,
+        region: string,
+        app: string,
+        useDeviceCodeFlow?: boolean
+    ): Promise<AuthError | undefined>
 
     async getAuthenticatedCredentialsError(data: StaticProfile): Promise<StaticProfileKeyErrorMessage | undefined> {
         return Auth.instance.authenticateData(data)
diff --git a/packages/core/src/login/webview/vue/login.vue b/packages/core/src/login/webview/vue/login.vue
@@ -208,6 +208,19 @@
                         {{ `${region.name} (${region.id})` }}
                     </option>
                 </select>
+                <div class="form-group topMargin">
+                    <div class="toggle-container">
+                        <label class="toggle-label">
+                            <input
+                                type="checkbox"
+                                id="useDeviceCodeFlow"
+                                name="useDeviceCodeFlow"
+                                v-model="useDeviceCodeFlow"
+                            />
+                            <span class="toggle-text">Use device code flow</span>
+                        </label>
+                    </div>
+                </div>
                 <button
                     class="continue-button topMargin"
                     :disabled="shouldDisableSsoContinue()"
@@ -350,6 +363,7 @@ export default defineComponent({
             profileName: '',
             accessKey: '',
             secretKey: '',
+            useDeviceCodeFlow: false,
         }
     },
     async created() {
@@ -447,7 +461,12 @@ export default defineComponent({
                 }
                 this.stage = 'AUTHENTICATING'
 
-                const error = await client.startEnterpriseSetup(this.startUrl, this.selectedRegion, this.app)
+                const error = await client.startEnterpriseSetup(
+                    this.startUrl,
+                    this.selectedRegion,
+                    this.app,
+                    this.useDeviceCodeFlow
+                )
                 if (error) {
                     this.stage = 'START'
                     void client.errorNotification(error)
@@ -833,8 +852,21 @@ body.vscode-light #logo-text {
     width: 10px;
 }
 
-.help-link__label {
-    margin: 0;
-    padding: 0 0 0 2px;
+.toggle-container {
+    margin-bottom: 10px;
+}
+
+.toggle-label {
+    display: flex;
+    align-items: center;
+    cursor: pointer;
+}
+
+.toggle-label input[type='checkbox'] {
+    margin-right: 8px;
+}
+
+.toggle-text {
+    font-size: var(--font-size-sm);
 }
 </style>
diff --git a/packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts b/packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts
@@ -40,13 +40,26 @@ export class ToolkitLoginWebview extends CommonAuthWebview {
         this.isCodeCatalystLogin = serviceToShow === 'codecatalyst'
     }
 
-    async startEnterpriseSetup(startUrl: string, region: string): Promise<AuthError | undefined> {
-        getLogger().debug(`called startEnterpriseSetup() with startUrl: '${startUrl}', region: '${region}'`)
+    async startEnterpriseSetup(
+        startUrl: string,
+        region: string,
+        app: string,
+        useDeviceCodeFlow?: boolean
+    ): Promise<AuthError | undefined> {
+        getLogger().debug(
+            `called startEnterpriseSetup() with startUrl: '${startUrl}', region: '${region}', useDeviceCodeFlow: ${useDeviceCodeFlow}`
+        )
         const metadata: TelemetryMetadata = {
             credentialSourceId: 'iamIdentityCenter',
             credentialStartUrl: startUrl,
             isReAuth: false,
         }
+
+        // Store the device code flow preference if it's set
+        if (useDeviceCodeFlow) {
+            await globals.globalState.update('aws.forceDeviceCodeFlow', true)
+        }
+
         await globals.globalState.update('recentSso', {
             startUrl: startUrl,
             region: region,
diff --git a/packages/core/src/shared/globalState.ts b/packages/core/src/shared/globalState.ts
@@ -50,6 +50,7 @@ export type globalKey =
     | 'aws.amazonq.customization.overrideV2'
     | 'aws.amazonq.regionProfiles'
     | 'aws.amazonq.regionProfiles.cache'
+    | 'aws.forceDeviceCodeFlow' // Toggle to force device code flow for authentication
     // Deprecated/legacy names. New keys should start with "aws.".
     | '#sessionCreationDates' // Legacy name from `ssoAccessTokenProvider.ts`.
     | 'CODECATALYST_RECONNECT'
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
@@ -62,7 +62,7 @@ describe('Amazon Q Login', function () {
     })
 
     it('signs into IdC and emits telemetry', async function () {
-        await backend.startEnterpriseSetup(startUrl, region)
+        await backend.startEnterpriseSetup(startUrl, region, 'AMAZONQ')
 
         assert.ok(isIdcSsoConnection(auth.activeConnection))
         assert.deepStrictEqual(auth.activeConnection.scopes, amazonQScopes)
diff --git a/packages/core/src/test/login/webview/vue/backend_toolkit.test.ts b/packages/core/src/test/login/webview/vue/backend_toolkit.test.ts
@@ -68,7 +68,7 @@ describe('Toolkit Login', function () {
     })
 
     it('signs into account IdC and emits telemetry', async function () {
-        await backend.startEnterpriseSetup(startUrl, region)
+        await backend.startEnterpriseSetup(startUrl, region, 'TOOLKIT')
 
         assert.ok(isIdcSsoConnection(auth.activeConnection))
         assert.deepStrictEqual(auth.activeConnection.scopes, scopesSsoAccountAccess)
@@ -90,7 +90,7 @@ describe('Toolkit Login', function () {
         sandbox.stub(codecatalystAuth, 'isConnectionOnboarded').resolves(true)
         backend.setLoginService('codecatalyst')
 
-        await backend.startEnterpriseSetup(startUrl, region)
+        await backend.startEnterpriseSetup(startUrl, region, 'TOOLKIT')
 
         assert.ok(isIdcSsoConnection(auth.activeConnection))
         assert.deepStrictEqual(auth.activeConnection.scopes, defaultScopes)
