Skip to content

Commit

Permalink
Wait for authentication to complete before you can save a cell (#1785)
Browse files Browse the repository at this point in the history
* Wait for authentication to complete before you can save a cell

* Words

---------

Co-authored-by: Sebastian (Tiedtke) Huckleberry <sebastiantiedtke@gmail.com>
  • Loading branch information
pastuxso and sourishkrout authored Jan 28, 2025
1 parent 17a6459 commit 7109023
Show file tree
Hide file tree
Showing 4 changed files with 160 additions and 80 deletions.
5 changes: 5 additions & 0 deletions __mocks__/vscode.ts
Original file line number Diff line number Diff line change
Expand Up @@ -327,4 +327,9 @@ export enum EndOfLine {
CRLF = 2
}

export class CancellationTokenSource {
cancel = vi.fn()
dispose = vi.fn()
}

export const version = '9.9.9'
89 changes: 81 additions & 8 deletions src/extension/messages/platformRequest/saveCellExecution.ts
Original file line number Diff line number Diff line change
@@ -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'
Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand All @@ -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()
Expand Down Expand Up @@ -318,5 +386,10 @@ export default async function saveCellExecution(
id: message.output.id,
hasErrors: true,
})
} finally {
if (currentCts?.token === token) {
currentCts.dispose()
currentCts = undefined
}
}
}
76 changes: 4 additions & 72 deletions src/extension/provider/statefulAuth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ import {
window,
AuthenticationSession,
AuthenticationProviderAuthenticationSessionsChangeEvent,
Event,
workspace,
AuthenticationGetSessionOptions,
} from 'vscode'
Expand All @@ -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
Expand All @@ -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 {
Expand Down Expand Up @@ -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) => {
Expand Down Expand Up @@ -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()
Expand Down
70 changes: 70 additions & 0 deletions src/utils/promiseFromEvent.ts
Original file line number Diff line number Diff line change
@@ -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,
}
}

0 comments on commit 7109023

Please sign in to comment.