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

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,117 +49,136 @@ export default class TabnineAuthenticationProvider
 
   private initializedDisposable: Disposable | undefined;
 
-  private lastState: State | undefined | null;
-
   private onDidLogin = new Emitter();
 
-  private myOnDidChangeSessions = new EventEmitter<AuthenticationProviderAuthenticationSessionsChangeEvent>();
+  private sessionsChangeEventEmitter = new EventEmitter<AuthenticationProviderAuthenticationSessionsChangeEvent>();
 
-  get onDidChangeSessions(): Event<AuthenticationProviderAuthenticationSessionsChangeEvent> {
-    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<AuthenticationProviderAuthenticationSessionsChangeEvent> {
+    return this.sessionsChangeEventEmitter.event;
+  }
+
   getSessions(): Promise<readonly AuthenticationSession[]> {
-    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<AuthenticationSession> {
     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<UserAuthData> {
+    return ((await once(this.onDidLogin, LOGIN_HAPPENED_EVENT)) as [
+      UserAuthData
+    ])[0];
   }
 
+  // eslint-disable-next-line class-methods-use-this
   async removeSession(): Promise<void> {
     await callForLogout();
 
-    this.myOnDidChangeSessions.fire({
-      removed: [(await this.getSessions())[0]],
-    });
     await sleep(5000);
   }
 
   dispose(): void {
     this.initializedDisposable?.dispose();
   }
 
-  private async waitForLogin(): Promise<State> {
-    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<void> {
-    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
   );
 }
 

src/binary/binaryStateSingleton.ts

@@ -25,7 +25,7 @@ export class BinaryState extends EventEmitterBasedState<State> {
     });
   }
 
-  private async checkForUpdates() {
+  async checkForUpdates() {
     try {
       await this.asyncSet(getStateOrNull);
     } catch (error) {

src/enterprise/tabnineChatWidget/SelfHostedChatEnabledState.ts

@@ -1,23 +1,27 @@
-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<ChatEnabledStateData>
   implements ChatEnabledState {
   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();
+        }
+      )
     );
   }
 

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();

src/state/EventEmitterBasedState.ts

@@ -36,7 +36,7 @@ export default class EventEmitterBasedState<T> implements Disposable {
     });
   }
 
-  onChange(subscription: (newValue: T) => void): Disposable {
+  onChange(subscription: (newValue: T) => unknown): Disposable {
     if (this.value !== null) {
       subscription(this.value);
     }

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<T> = Disposable & EventEmitterBasedState<T>;
+export type DerivedNonNullState<T> = Disposable &
+  EventEmitterBasedNonNullState<T>;
 
-function deriveState<I, O, S extends EventEmitterBasedState<I>>(
-  state: S,
+function deriveState<I, O>(
+  state: EventEmitterBasedState<I>,
   mapping: (value: I) => O
 ): DerivedState<O> {
-  class TempDerivedState
-    extends EventEmitterBasedState<O>
-    implements Disposable {
+  return new (class extends EventEmitterBasedState<O> implements Disposable {
     useStateDisposabled!: Disposable;
 
     constructor() {
@@ -24,13 +26,36 @@ function deriveState<I, O, S extends EventEmitterBasedState<I>>(
       super.dispose();
       this.useStateDisposabled.dispose();
     }
-  }
+  })();
+}
+
+export function deriveNonNullState<I, O>(
+  state: EventEmitterBasedState<I>,
+  mapping: (value: I, self: O) => O,
+  initailValue: O
+): DerivedNonNullState<O> {
+  return new (class
+    extends EventEmitterBasedNonNullState<O>
+    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<I, O, S extends EventEmitterBasedState<I>>(
-  state: S,
+export function useDerviedState<I, O>(
+  state: EventEmitterBasedState<I>,
   mapping: (value: I) => O,
   onChange: (newValue: O) => void
 ): Disposable {