Fix authentication stuck after logout

codota · Nov 28, 2023 · 0274793 · 0274793
1 parent 8b5a67f
commit 0274793
Show file tree

Hide file tree

Showing 8 changed files with 150 additions and 82 deletions.
diff --git 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 "../utils/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.useState((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
   );
 }
 

diff --git a/src/binary/binaryStateSingleton.ts b/src/binary/binaryStateSingleton.ts
@@ -31,7 +31,7 @@ export class BinaryState extends EventEmitterBasedState<State> {
     return this.intervalDisposabled;
   }
 
-  private async checkForUpdates() {
+  async checkForUpdates() {
     try {
       await this.updateStateLock.use(async () => {
         const state = await getState();

diff --git a/src/extension.ts b/src/extension.ts
@@ -121,9 +121,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/tabnineChatWidget/ChatEnabledState.ts b/src/tabnineChatWidget/ChatEnabledState.ts
@@ -30,5 +30,5 @@ export const ChatStates = {
 };
 
 export default interface ChatEnabledState {
-  useState(onChange: (state: ChatEnabledStateData) => void): Disposable;
+  onChange(subscription: (state: ChatEnabledStateData) => void): Disposable;
 }
diff --git a/src/tabnineChatWidget/SaasChatEnabledState.ts b/src/tabnineChatWidget/SaasChatEnabledState.ts
@@ -31,13 +31,13 @@ export default class SaasChatEnabledState
         }
       ),
       onDidRefreshCapabilities(() => {
-        this.updateState(BINARY_STATE.get()?.is_logged_in || false);
+        this.updateState(BINARY_STATE.get()?.is_logged_in ?? null);
       })
     );
   }
 
-  private updateState(isLoggedIn: boolean) {
-    if (!isCapabilitiesReady()) {
+  private updateState(isLoggedIn: boolean | null) {
+    if (!isCapabilitiesReady() || isLoggedIn == null) {
       return;
     }
 

diff --git a/src/tabnineChatWidget/tabnineChatWidgetWebview.ts b/src/tabnineChatWidget/tabnineChatWidgetWebview.ts
@@ -25,7 +25,7 @@ export default function registerTabnineChatWidgetWebview(
   setTabnineChatWebview("loading");
 
   context.subscriptions.push(
-    chatEnabledState.useState((state) => {
+    chatEnabledState.onChange((state) => {
       if (state.enabled) {
         registerChatView(serverUrl, context);
       } else if (state.chatNotEnabledReason) {

diff --git a/src/utils/EventEmitterBasedState.ts b/src/utils/EventEmitterBasedState.ts
@@ -23,11 +23,11 @@ export default class EventEmitterBasedState<T> {
     }
   }
 
-  useState(onChange: (newValue: T) => void): Disposable {
+  onChange(subscription: (newValue: T) => unknown): Disposable {
     if (this.value !== null) {
-      onChange(this.value);
+      subscription(this.value);
     }
 
-    return this.eventEmitter.event(onChange);
+    return this.eventEmitter.event(subscription);
   }
 }
diff --git a/src/utils/deriveState.ts b/src/utils/deriveState.ts
@@ -1,7 +1,11 @@
+// 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>;
 
 export default function deriveState<I, O, S extends EventEmitterBasedState<I>>(
   state: S,
@@ -15,7 +19,7 @@ export default function deriveState<I, O, S extends EventEmitterBasedState<I>>(
     constructor() {
       super();
 
-      this.useStateDisposabled = state.useState((inputState) => {
+      this.useStateDisposabled = state.onChange((inputState) => {
         this.set(mapping(inputState));
       });
     }
@@ -28,13 +32,39 @@ export default function deriveState<I, O, S extends EventEmitterBasedState<I>>(
   return new TempDerivedState();
 }
 
+export function deriveNonNullState<I, O, S extends EventEmitterBasedState<I>>(
+  state: S,
+  mapping: (value: I, self: O) => O,
+  initailValue: O
+): DerivedNonNullState<O> {
+  class TempDerivedNonNullState
+    extends EventEmitterBasedNonNullState<O>
+    implements Disposable {
+    useStateDisposabled!: Disposable;
+
+    constructor() {
+      super(initailValue);
+
+      this.useStateDisposabled = state.onChange((inputState) => {
+        this.set(mapping(inputState, this.get()));
+      });
+    }
+
+    dispose() {
+      this.useStateDisposabled.dispose();
+    }
+  }
+
+  return new TempDerivedNonNullState();
+}
+
 export function useDerviedState<I, O, S extends EventEmitterBasedState<I>>(
   state: S,
   mapping: (value: I) => O,
   onChange: (newValue: O) => void
 ): Disposable {
   const derviedState = deriveState(state, mapping);
-  const disposable = derviedState.useState(onChange);
+  const disposable = derviedState.onChange(onChange);
 
   return {
     dispose() {