Wait for authentication to complete before you can save a cell

__mocks__/vscode.ts

@@ -327,4 +327,9 @@ export enum EndOfLine {
   CRLF = 2
 }
 
+export class CancellationTokenSource {
+  cancel = vi.fn()
+  dispose = vi.fn()
+}
+
 export const version = '9.9.9'

src/extension/messages/platformRequest/saveCellExecution.ts

@@ -1,6 +1,15 @@
 import os from 'node:os'
 
-import { Uri, env, workspace, commands } from 'vscode'
+import {
+  Uri,
+  env,
+  workspace,
+  commands,
+  EventEmitter,
+  AuthenticationSessionsChangeEvent,
+  window,
+  CancellationTokenSource,
+} from 'vscode'
 import { TelemetryReporter } from 'vscode-telemetry'
 import getMAC from 'getmac'
 import YAML from 'yaml'
@@ -31,13 +40,22 @@ import {
 } from '../../__generated-platform__/graphql'
 import { Frontmatter } from '../../grpc/parser/tcp/types'
 import { getCellById } from '../../cell'
-import { StatefulAuthProvider } from '../../provider/statefulAuth'
+import {
+  AUTH_TIMEOUT,
+  StatefulAuthProvider,
+  StatefulAuthSession,
+} from '../../provider/statefulAuth'
 import features from '../../features'
+import AuthSessionChangeHandler from '../../authSessionChangeHandler'
+import { promiseFromEvent } from '../../../utils/promiseFromEvent'
 import { getDocumentCacheId } from '../../serializer/serializer'
 import { ConnectSerializer } from '../../serializer'
 export type APIRequestMessage = IApiMessage<ClientMessage<ClientMessages.platformApiRequest>>
 
 const log = getLogger('SaveCell')
+type SessionType = StatefulAuthSession | undefined
+
+let currentCts: CancellationTokenSource | undefined
 
 export default async function saveCellExecution(
   requestMessage: APIRequestMessage,
@@ -46,6 +64,13 @@ export default async function saveCellExecution(
   const isReporterEnabled = features.isOnInContextState(FeatureName.ReporterAPI)
   const { messaging, message, editor } = requestMessage
 
+  if (currentCts) {
+    currentCts.cancel()
+  }
+
+  currentCts = new CancellationTokenSource()
+  const { token } = currentCts
+
   try {
     const autoSaveIsOn = ContextState.getKey<boolean>(NOTEBOOK_AUTOSAVE_ON)
     const forceLogin = kernel.isFeatureOn(FeatureName.ForceLogin)
@@ -58,12 +83,55 @@ export default async function saveCellExecution(
 
     if (!session && message.output.data.isUserAction) {
       await commands.executeCommand('runme.openCloudPanel')
-      return postClientMessage(messaging, ClientMessages.platformApiResponse, {
-        data: {
-          displayShare: false,
-        },
-        id: message.output.id,
-      })
+
+      const authenticationEvent = new EventEmitter<StatefulAuthSession | undefined>()
+
+      const callback = (_e: AuthenticationSessionsChangeEvent) => {
+        AuthSessionChangeHandler.instance.removeListener(callback)
+        StatefulAuthProvider.instance.currentSession().then((session) => {
+          authenticationEvent.fire(session)
+        })
+      }
+
+      AuthSessionChangeHandler.instance.addListener(callback)
+
+      if (token.isCancellationRequested) {
+        return
+      }
+
+      try {
+        session = await Promise.race([
+          promiseFromEvent<SessionType, SessionType>(authenticationEvent.event).promise,
+          new Promise<undefined>((resolve, reject) => {
+            const timeoutId = setTimeout(() => reject(undefined), AUTH_TIMEOUT)
+            token.onCancellationRequested(() => {
+              clearTimeout(timeoutId)
+              reject(new Error('Operation cancelled'))
+            })
+          }),
+        ])
+      } finally {
+        authenticationEvent.dispose()
+
+        if (token.isCancellationRequested) {
+          log.info('Cancelling authentication event')
+          return
+        }
+
+        if (!session) {
+          await postClientMessage(messaging, ClientMessages.platformApiResponse, {
+            data: {
+              displayShare: false,
+            },
+            id: message.output.id,
+          })
+
+          window.showWarningMessage(
+            'Saving timed out. Sign in to save your cells. Please try again.',
+          )
+          return
+        }
+      }
     }
 
     const graphClient = await InitializeCloudClient()
@@ -318,5 +386,10 @@ export default async function saveCellExecution(
       id: message.output.id,
       hasErrors: true,
     })
+  } finally {
+    if (currentCts?.token === token) {
+      currentCts.dispose()
+      currentCts = undefined
+    }
   }
 }

src/extension/provider/statefulAuth.ts

@@ -12,7 +12,6 @@ import {
   window,
   AuthenticationSession,
   AuthenticationProviderAuthenticationSessionsChangeEvent,
-  Event,
   workspace,
   AuthenticationGetSessionOptions,
 } from 'vscode'
@@ -31,12 +30,14 @@ import ContextState from '../contextState'
 import getLogger from '../logger'
 import { FeatureName } from '../../types'
 import * as features from '../features'
+import { PromiseAdapter, promiseFromEvent } from '../../utils/promiseFromEvent'
 
 const logger = getLogger('StatefulAuthProvider')
 
 const AUTH_NAME = 'Stateful'
 const SESSIONS_SECRET_KEY = `${AuthenticationProviders.Stateful}.sessions`
 export const DEFAULT_SCOPES = ['profile']
+export const AUTH_TIMEOUT = 60000
 
 interface TokenInformation {
   accessToken: string
@@ -53,21 +54,6 @@ interface DecodedToken extends JwtPayload {
   scope?: string
 }
 
-// Interface declaration for a PromiseAdapter
-interface PromiseAdapter<T, U> {
-  // Function signature of the PromiseAdapter
-  (
-    // Input value of type T that the adapter function will process
-    value: T,
-    // Function to resolve the promise with a value of type U or a promise that resolves to type U
-    resolve: (value: U | PromiseLike<U>) => void,
-    // Function to reject the promise with a reason of any type
-    reject: (reason: any) => void,
-  ): any // The function can return a value of any type
-}
-
-const passthrough = (value: any, resolve: (value?: any) => void) => resolve(value)
-
 type SessionsChangeEvent = AuthenticationProviderAuthenticationSessionsChangeEvent
 
 export class StatefulAuthProvider implements AuthenticationProvider, Disposable {
@@ -478,8 +464,8 @@ export class StatefulAuthProvider implements AuthenticationProvider, Disposable
           return await Promise.race([
             // Waiting for the codeExchangePromise to resolve
             codeExchangePromise.promise,
-            // Creating a new promise that rejects after 60000 milliseconds
-            new Promise<string>((_, reject) => setTimeout(() => reject('Cancelled'), 60000)),
+            // Creating a new promise that rejects on timeout
+            new Promise<string>((_, reject) => setTimeout(() => reject('Cancelled'), AUTH_TIMEOUT)),
             // Creating a promise based on an event, rejecting with 'User Cancelled' when
             // token.onCancellationRequested event occurs
             promiseFromEvent<any, any>(token.onCancellationRequested, (_, __, reject) => {
@@ -715,60 +701,6 @@ export class StatefulAuthProvider implements AuthenticationProvider, Disposable
   }
 }
 
-/**
- * Return a promise that resolves with the next emitted event, or with some future
- * event as decided by an adapter.
- *
- * If specified, the adapter is a function that will be called with
- * `(event, resolve, reject)`. It will be called once per event until it resolves or
- * rejects.
- *
- * The default adapter is the passthrough function `(value, resolve) => resolve(value)`.
- *
- * @param event the event
- * @param adapter controls resolution of the returned promise
- * @returns a promise that resolves or rejects as specified by the adapter
- */
-function promiseFromEvent<T, U>(
-  event: Event<T>,
-  adapter: PromiseAdapter<T, U> = passthrough,
-): { promise: Promise<U>; cancel: EventEmitter<void> } {
-  let subscription: Disposable
-  let cancel = new EventEmitter<void>()
-
-  // Return an object containing a promise and a cancel EventEmitter
-  return {
-    // Creating a new Promise
-    promise: new Promise<U>((resolve, reject) => {
-      // Listening for the cancel event and rejecting the promise with 'Cancelled' when it occurs
-      cancel.event((_) => reject('Cancelled'))
-      // Subscribing to the event
-      subscription = event((value: T) => {
-        try {
-          // Resolving the promise with the result of the adapter function
-          Promise.resolve(adapter(value, resolve, reject)).catch(reject)
-        } catch (error) {
-          // Rejecting the promise if an error occurs during execution
-          reject(error)
-        }
-      })
-    }).then(
-      // Disposing the subscription and returning the result when the promise resolves
-      (result: U) => {
-        subscription.dispose()
-        return result
-      },
-      // Disposing the subscription and re-throwing the error when the promise rejects
-      (error) => {
-        subscription.dispose()
-        throw error
-      },
-    ),
-    // Returning the cancel EventEmitter
-    cancel,
-  }
-}
-
 function secsToUnixTime(seconds: number) {
   const now = new Date()
   return new Date(now.getTime() + seconds * 1000).getTime()

src/utils/promiseFromEvent.ts

@@ -0,0 +1,70 @@
+import { Disposable, Event, EventEmitter } from 'vscode'
+
+// Interface declaration for a PromiseAdapter
+export interface PromiseAdapter<T, U> {
+  // Function signature of the PromiseAdapter
+  (
+    // Input value of type T that the adapter function will process
+    value: T,
+    // Function to resolve the promise with a value of type U or a promise that resolves to type U
+    resolve: (value: U | PromiseLike<U>) => void,
+    // Function to reject the promise with a reason of any type
+    reject: (reason: any) => void,
+  ): any // The function can return a value of any type
+}
+
+const passthrough = (value: any, resolve: (value?: any) => void) => resolve(value)
+
+/**
+ * Return a promise that resolves with the next emitted event, or with some future
+ * event as decided by an adapter.
+ *
+ * If specified, the adapter is a function that will be called with
+ * `(event, resolve, reject)`. It will be called once per event until it resolves or
+ * rejects.
+ *
+ * The default adapter is the passthrough function `(value, resolve) => resolve(value)`.
+ *
+ * @param event the event
+ * @param adapter controls resolution of the returned promise
+ * @returns a promise that resolves or rejects as specified by the adapter
+ */
+export function promiseFromEvent<T, U>(
+  event: Event<T>,
+  adapter: PromiseAdapter<T, U> = passthrough,
+): { promise: Promise<U>; cancel: EventEmitter<void> } {
+  let subscription: Disposable
+  let cancel = new EventEmitter<void>()
+
+  // Return an object containing a promise and a cancel EventEmitter
+  return {
+    // Creating a new Promise
+    promise: new Promise<U>((resolve, reject) => {
+      // Listening for the cancel event and rejecting the promise with 'Cancelled' when it occurs
+      cancel.event((_) => reject('Cancelled'))
+      // Subscribing to the event
+      subscription = event((value: T) => {
+        try {
+          // Resolving the promise with the result of the adapter function
+          Promise.resolve(adapter(value, resolve, reject)).catch(reject)
+        } catch (error) {
+          // Rejecting the promise if an error occurs during execution
+          reject(error)
+        }
+      })
+    }).then(
+      // Disposing the subscription and returning the result when the promise resolves
+      (result: U) => {
+        subscription.dispose()
+        return result
+      },
+      // Disposing the subscription and re-throwing the error when the promise rejects
+      (error) => {
+        subscription.dispose()
+        throw error
+      },
+    ),
+    // Returning the cancel EventEmitter
+    cancel,
+  }
+}