From c401753a93bf3042732a94f77e041cdea773bf4e Mon Sep 17 00:00:00 2001 From: ofekby <39569540+ofekby@users.noreply.github.com> Date: Mon, 4 Dec 2023 18:07:55 +0200 Subject: [PATCH] DEV2-4362: Fix sign-in not available after sign-out (#1381) * Chagne to use state to infer showing join a team * Change chat to use state to show views * Fix authentication stuck after logout * Remove unused generics * Move lock to the state management * Fix sh state too * Renaming * Remove duplicate of binary state * Remove unused export * Fix linting issues * Fix exports * Fix disposable * Make the classes anonymous * TEMP commit * Change to just listen to binary state --- .../TabnineAuthenticationProvider.ts | 177 +++++++++++------- src/binary/binaryStateSingleton.ts | 2 +- .../SelfHostedChatEnabledState.ts | 16 +- src/extension.ts | 3 - src/state/EventEmitterBasedState.ts | 2 +- src/state/deriveState.ts | 43 ++++- 6 files changed, 155 insertions(+), 88 deletions(-) diff --git a/src/authentication/TabnineAuthenticationProvider.ts b/src/authentication/TabnineAuthenticationProvider.ts index 5e83d18627..2cb61e3fce 100644 --- a/src/authentication/TabnineAuthenticationProvider.ts +++ b/src/authentication/TabnineAuthenticationProvider.ts @@ -9,16 +9,38 @@ import { EventEmitter, } from "vscode"; import { once, EventEmitter as Emitter } from "events"; +import { once as runOnce } from "underscore"; import { State } from "../binary/state"; import { BRAND_NAME } from "../globals/consts"; import { sleep } from "../utils/utils"; import { callForLogin, callForLogout } from "./authentication.api"; import TabnineSession from "./TabnineSession"; import BINARY_STATE from "../binary/binaryStateSingleton"; -import { getState } from "../binary/requests/requests"; +import { deriveNonNullState } from "../state/deriveState"; const LOGIN_HAPPENED_EVENT = "loginHappened"; +type UserAuthData = { + username: string; + accessToken: string; +}; + +type AuthStateData = { + current: UserAuthData | null; + last: UserAuthData | null; +}; + +const AUTH_INITIAL_STATE = { + last: null, + current: null, +}; + +function toSession({ accessToken, username }: UserAuthData): TabnineSession { + return new TabnineSession(username, accessToken); +} + +const setAuthenticationReadyOnce = runOnce(setAuthenticationReady); + export default class TabnineAuthenticationProvider implements AuthenticationProvider, Disposable { public readonly id: string = BRAND_NAME; @@ -27,45 +49,51 @@ export default class TabnineAuthenticationProvider private initializedDisposable: Disposable | undefined; - private lastState: State | undefined | null; - private onDidLogin = new Emitter(); - private myOnDidChangeSessions = new EventEmitter(); + private sessionsChangeEventEmitter = new EventEmitter(); - get onDidChangeSessions(): Event { - return this.myOnDidChangeSessions.event; - } + private authState = deriveNonNullState( + BINARY_STATE, + calculateAuthState, + AUTH_INITIAL_STATE + ); constructor() { this.initializedDisposable = Disposable.from( - this.handleSessionChange(), - this.pollState() + this.authState, + this.onDerivedAuthStateChanged(), + listenForSessionChangeFromVscode() ); } + get onDidChangeSessions(): Event { + return this.sessionsChangeEventEmitter.event; + } + getSessions(): Promise { - const state = this.lastState; + const userData = this.authState.get().current; - return Promise.resolve( - state?.is_logged_in - ? [new TabnineSession(state?.user_name, state?.access_token)] - : [] - ); + return Promise.resolve(userData ? [toSession(userData)] : []); } async createSession(): Promise { await callForLogin(); - const state = await this.waitForLogin(); - return new TabnineSession(state?.user_name, state?.access_token); + const userAuth = await this.waitForLogin(); + + return toSession(userAuth); + } + + private async waitForLogin(): Promise { + return ((await once(this.onDidLogin, LOGIN_HAPPENED_EVENT)) as [ + UserAuthData + ])[0]; } + // eslint-disable-next-line class-methods-use-this async removeSession(): Promise { await callForLogout(); - this.myOnDidChangeSessions.fire({ - removed: [(await this.getSessions())[0]], - }); await sleep(5000); } @@ -73,71 +101,84 @@ export default class TabnineAuthenticationProvider this.initializedDisposable?.dispose(); } - private async waitForLogin(): Promise { - return ((await once(this.onDidLogin, LOGIN_HAPPENED_EVENT)) as [State])[0]; - } + private onDerivedAuthStateChanged(): Disposable { + return this.authState.onChange(async ({ current, last }) => { + await setAuthenticationReadyOnce(); - private handleSessionChange(): Disposable { - // This fires when the user initiates a "silent" auth flow via the Accounts menu. - return authentication.onDidChangeSessions((e) => { - if (e.provider.id === BRAND_NAME) { - void getState().then((state) => { - void this.checkForUpdates(state); - }); + if (current && !last) { + this.onDidLogin.emit(LOGIN_HAPPENED_EVENT, current); + } + + if (!current) { + await clearSessionPreference(); } - }); - } - private pollState(): Disposable { - return BINARY_STATE.onChange((state) => { - void this.checkForUpdates(state); + await setAuthenticationState(Boolean(current)); + this.notifyVscodeOfAuthStateChanges(current, last); }); } - private async checkForUpdates( - state: State | null | undefined - ): Promise { - const added: AuthenticationSession[] = []; - const removed: AuthenticationSession[] = []; - - const { lastState } = this; - - this.lastState = state; - - const newState = this.lastState; - const oldState = lastState; + private notifyVscodeOfAuthStateChanges( + current: UserAuthData | null, + last: UserAuthData | null + ) { + if (!last && current) { + this.sessionsChangeEventEmitter.fire({ + added: [toSession(current)], + }); + } - if (newState?.is_logged_in) { - this.onDidLogin.emit(LOGIN_HAPPENED_EVENT, newState); + if (last && !current) { + this.sessionsChangeEventEmitter.fire({ + removed: [toSession(last)], + }); } - if (newState) { - await setAuthenticationReady(); + if (last && current) { + this.sessionsChangeEventEmitter.fire({ + removed: [toSession(last)], + added: [toSession(current)], + }); } - await setAuthenticationState(oldState, newState); - - if (!oldState?.is_logged_in && newState?.is_logged_in) { - added.push((await this.getSessions())[0]); - } else if (newState && !newState.is_logged_in && oldState?.is_logged_in) { - removed.push((await this.getSessions())[0]); - } else { - return; + } +} + +async function clearSessionPreference() { + await authentication.getSession(BRAND_NAME, [], { + clearSessionPreference: true, + }); +} + +function listenForSessionChangeFromVscode(): Disposable { + // This fires when the user initiates a "silent" auth flow via the Accounts menu. + return authentication.onDidChangeSessions((e) => { + if (e.provider.id === BRAND_NAME) { + void BINARY_STATE.checkForUpdates(); } + }); +} - this.myOnDidChangeSessions.fire({ - added, - removed, - }); +function calculateAuthState(binartState: State, value: AuthStateData) { + const newValue: AuthStateData = { + last: value.current, + current: null, + }; + + if (binartState.is_logged_in) { + newValue.current = { + accessToken: binartState.access_token || "", + username: binartState.user_name, + }; } + + return newValue; } -async function setAuthenticationState( - oldState: State | null | undefined, - newState: State | null | undefined -) { + +async function setAuthenticationState(authenticated: boolean) { return commands.executeCommand( "setContext", "tabnine.authenticated", - oldState?.is_logged_in || newState?.is_logged_in + authenticated ); } diff --git a/src/binary/binaryStateSingleton.ts b/src/binary/binaryStateSingleton.ts index f393961d45..a2b3fc867e 100644 --- a/src/binary/binaryStateSingleton.ts +++ b/src/binary/binaryStateSingleton.ts @@ -25,7 +25,7 @@ export class BinaryState extends EventEmitterBasedState { }); } - private async checkForUpdates() { + async checkForUpdates() { try { await this.asyncSet(getStateOrNull); } catch (error) { diff --git a/src/enterprise/tabnineChatWidget/SelfHostedChatEnabledState.ts b/src/enterprise/tabnineChatWidget/SelfHostedChatEnabledState.ts index 15cebb881c..7784191a05 100644 --- a/src/enterprise/tabnineChatWidget/SelfHostedChatEnabledState.ts +++ b/src/enterprise/tabnineChatWidget/SelfHostedChatEnabledState.ts @@ -1,10 +1,12 @@ -import { ExtensionContext, authentication } from "vscode"; +import { ExtensionContext } from "vscode"; import ChatEnabledState, { ChatEnabledStateData, ChatStates, } from "../../tabnineChatWidget/ChatEnabledState"; import EventEmitterBasedNonNullState from "../../state/EventEmitterBasedNonNullState"; import getUserInfo from "../requests/UserInfo"; +import { useDerviedState } from "../../state/deriveState"; +import BINARY_STATE from "../../binary/binaryStateSingleton"; export default class SelfHostedChatEnabledState extends EventEmitterBasedNonNullState @@ -12,12 +14,14 @@ export default class SelfHostedChatEnabledState constructor(context: ExtensionContext) { super(ChatStates.loading); - void this.updateState(); - context.subscriptions.push( - authentication.onDidChangeSessions(() => { - void this.updateState(); - }) + useDerviedState( + BINARY_STATE, + (state) => state.is_logged_in, + () => { + void this.updateState(); + } + ) ); } diff --git a/src/extension.ts b/src/extension.ts index b677b329e0..5fa65cf994 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -120,9 +120,6 @@ async function backgroundInit(context: vscode.ExtensionContext) { new TabnineAuthenticationProvider() ) ); - await vscode.authentication.getSession(BRAND_NAME, [], { - clearSessionPreference: true, - }); } vscode.commands.registerCommand("tabnine.authenticate", () => { void callForLogin(); diff --git a/src/state/EventEmitterBasedState.ts b/src/state/EventEmitterBasedState.ts index d91630b588..fcf1efdfff 100644 --- a/src/state/EventEmitterBasedState.ts +++ b/src/state/EventEmitterBasedState.ts @@ -36,7 +36,7 @@ export default class EventEmitterBasedState implements Disposable { }); } - onChange(subscription: (newValue: T) => void): Disposable { + onChange(subscription: (newValue: T) => unknown): Disposable { if (this.value !== null) { subscription(this.value); } diff --git a/src/state/deriveState.ts b/src/state/deriveState.ts index 6ed85cd4db..f3231f3822 100644 --- a/src/state/deriveState.ts +++ b/src/state/deriveState.ts @@ -1,15 +1,17 @@ +// eslint-disable-next-line max-classes-per-file import { Disposable } from "vscode"; import EventEmitterBasedState from "./EventEmitterBasedState"; +import EventEmitterBasedNonNullState from "./EventEmitterBasedNonNullState"; export type DerivedState = Disposable & EventEmitterBasedState; +export type DerivedNonNullState = Disposable & + EventEmitterBasedNonNullState; -function deriveState>( - state: S, +function deriveState( + state: EventEmitterBasedState, mapping: (value: I) => O ): DerivedState { - class TempDerivedState - extends EventEmitterBasedState - implements Disposable { + return new (class extends EventEmitterBasedState implements Disposable { useStateDisposabled!: Disposable; constructor() { @@ -24,13 +26,36 @@ function deriveState>( super.dispose(); this.useStateDisposabled.dispose(); } - } + })(); +} + +export function deriveNonNullState( + state: EventEmitterBasedState, + mapping: (value: I, self: O) => O, + initailValue: O +): DerivedNonNullState { + return new (class + extends EventEmitterBasedNonNullState + implements Disposable { + useStateDisposabled!: Disposable; + + constructor() { + super(initailValue); - return new TempDerivedState(); + this.useStateDisposabled = state.onChange((inputState) => { + this.set(mapping(inputState, this.get())); + }); + } + + dispose() { + super.dispose(); + this.useStateDisposabled.dispose(); + } + })(); } -export function useDerviedState>( - state: S, +export function useDerviedState( + state: EventEmitterBasedState, mapping: (value: I) => O, onChange: (newValue: O) => void ): Disposable {