From f7a3b35ca825f7d9751757abd9205597f3df5b89 Mon Sep 17 00:00:00 2001
From: f1ames <dr.odpowiedz@gmail.com>
Date: Tue, 5 Dec 2023 14:39:01 +0100
Subject: [PATCH] fix: improve login flow

---
 src/commands/login.ts | 57 +++++++++++++++++++++++++++++++++++--------
 1 file changed, 47 insertions(+), 10 deletions(-)

diff --git a/src/commands/login.ts b/src/commands/login.ts
index 2b5f108..5d0fdd9 100644
--- a/src/commands/login.ts
+++ b/src/commands/login.ts
@@ -1,4 +1,4 @@
-import { env, Uri } from 'vscode';
+import { env, Uri, window } from 'vscode';
 import { canRun, disabledForLocal } from '../utils/commands';
 import { raiseError, raiseInfo } from '../utils/errors';
 import { COMMAND_NAMES } from '../constants';
@@ -7,8 +7,12 @@ import logger from '../utils/logger';
 import globals from '../utils/globals';
 import type { RuntimeContext } from '../utils/runtime-context';
 
+let pendingLoginRequest: Awaited<ReturnType<RuntimeContext['authenticator']['login']>> | undefined = undefined;
+
 export function getLoginCommand(context: RuntimeContext) {
   return async () => {
+    abortPendingLogin();
+
     if (!canRun() || disabledForLocal(context, COMMAND_NAMES.LOGIN)) {
       return;
     }
@@ -34,31 +38,45 @@ export function getLoginCommand(context: RuntimeContext) {
     const method = 'device code';
 
     try {
-      const loginRequest = await authenticator.login(method);
+      const loginRequest = pendingLoginRequest = await authenticator.login(method);
+
+      if (!loginRequest) {
+        abortPendingLogin();
+
+        trackEvent('command/login', {
+          status: 'cancelled',
+          method,
+        });
+
+        return;
+      }
+
       const handle =  loginRequest.handle;
 
-      raiseInfo(
+      const loginPromptResult = await window.showInformationMessage(
         `Please open ${handle.verification_uri_complete} and enter the code ${handle.user_code} to login.`,
-        [{
-          title: 'Open login page',
-          callback: () => {
-            env.openExternal(Uri.parse(handle.verification_uri_complete));
-          }
-        }],
         {
           modal: true,
         },
+        {title:'Open login page'}
       );
 
-      if (!loginRequest) {
+      if (!loginPromptResult) {
+        abortPendingLogin();
+
         trackEvent('command/login', {
           status: 'cancelled',
           method,
+          error: 'User cancelled login prompt.'
         });
 
         return;
       }
 
+      // We can't really rely on the return value of 'env.openExternal' since it returns false in any case but 'Open' button
+      // (e.g. for copy button). So this is not a good indication of whether the user wants to continue or not.
+      await env.openExternal(Uri.parse(handle.verification_uri_complete));
+
       const user = await loginRequest.onDone;
 
       raiseInfo(`You are now logged in as ${user.email}.`);
@@ -68,14 +86,33 @@ export function getLoginCommand(context: RuntimeContext) {
         method,
       });
     } catch (err) {
+      if (err.message === 'polling aborted') {
+        // This is expected upon cancellation of pending login request.
+        logger.log('Polling aborted', err);
+      } else {
         logger.error(err);
         raiseError(`Failed to login to Monokle Cloud. Please try again. Error: ${err.message}`);
 
+        abortPendingLogin();
+
         trackEvent('command/login', {
           status: 'failure',
           method,
           error: err.message,
         });
+      }
     }
   };
 }
+
+function abortPendingLogin() {
+  try {
+    if (pendingLoginRequest) {
+      const handle = pendingLoginRequest.handle;
+      pendingLoginRequest = undefined;
+      handle?.abort();
+    }
+  } catch (err) {
+    logger.error(err);
+  }
+}