DEV2-4157: Add tabnine chat state

package.json

@@ -136,10 +136,13 @@
     "@types/semver": "^7.3.4",
     "@types/sinon": "^9.0.11",
     "@types/tmp": "^0.2.0",
+    "@types/underscore": "^1.11.15",
     "@types/vscode": "^1.50.0",
     "@types/yauzl": "^2.9.1",
     "@typescript-eslint/eslint-plugin": "^5.45.0",
     "@typescript-eslint/parser": "^4.18.0",
+    "@vscode/test-electron": "^2.3.6",
+    "@vscode/vsce": "^2.22.0",
     "assert": "^2.0.0",
     "chai": "^4.2.0",
     "chai-shallow-deep-equal": "^1.4.6",
@@ -164,8 +167,6 @@
     "ts-loader": "^9.4.1",
     "ts-mockito": "^2.6.1",
     "typescript": "^4.2.2",
-    "@vscode/vsce": "^2.22.0",
-    "@vscode/test-electron": "^2.3.6",
     "webpack": "^5.75.0",
     "webpack-bundle-analyzer": "^4.7.0",
     "webpack-cli": "^5.0.0"
@@ -181,6 +182,7 @@
     "semver": "^7.3.2",
     "systeminformation": "^5.6.10",
     "tmp": "^0.2.1",
+    "underscore": "^1.13.6",
     "vscode-extension-telemetry": "^0.1.7"
   },
   "capabilities": {
@@ -217,24 +219,30 @@
           "type": "webview",
           "id": "tabnine.chat",
           "name": "Chat",
-          "when": "tabnine.authenticated && tabnine.chat.ready"
+          "when": "tabnine.chat.webview == 'chat'"
         },
         {
           "id": "tabnine.chat.welcome",
           "type": "webview",
           "name": "Welcome to Chat",
-          "when": "tabnine.authenticated && !tabnine.chat.ready"
+          "when": "tabnine.chat.webview == 'capability_required'"
         },
         {
-          "id": "tabnine.authenticate",
+          "id": "tabnine.chat.authenticate",
           "type": "webview",
           "name": "Please log in",
-          "when": "!tabnine.authenticated && tabnine.process.ready && (tabnine.enterprise || tabnine.capabilities.ready) && tabnine.authentication.ready"
+          "when": "tabnine.chat.webview == 'authnetication_required'"
+        },
+        {
+          "id": "tabnine.chat.not_part_of_a_team",
+          "type": "webview",
+          "name": "Please join a team",
+          "when": "tabnine.chat.webview == 'part_of_a_team_required'"
         },
         {
           "id": "tabnine.loading",
           "name": "Loading",
-          "when": "!tabnine.process.ready || !(tabnine.enterprise || tabnine.capabilities.ready) || !tabnine.authentication.ready"
+          "when": "tabnine.chat.webview == 'loading'"
         }
       ]
     },
@@ -262,17 +270,17 @@
       {
         "command": "tabnine.chat.submit-message",
         "title": "Submit message",
-        "when": "tabnine.chat.eval"
+        "when": "tabnine.chat.eval && tabnine.chat.ready"
       },
       {
         "command": "tabnine.chat.state",
         "title": "Chat global state",
-        "when": "tabnine.chat.eval"
+        "when": "tabnine.chat.eval && tabnine.chat.ready"
       },
       {
         "command": "tabnine.chat.clear-all-conversations",
         "title": "Clear all tabnine chat conversations",
-        "when": "tabnine.chat.eval"
+        "when": "tabnine.chat.eval && tabnine.chat.ready"
       },
       {
         "command": "TabNine::config",

src/authentication/TabnineAuthenticationProvider.ts

@@ -9,14 +9,14 @@ import {
   EventEmitter,
 } from "vscode";
 import { once, EventEmitter as Emitter } from "events";
-import { getState, tabNineProcess } from "../binary/requests/requests";
 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";
 
-const SESSION_POLL_INTERVAL = 10000;
 const LOGIN_HAPPENED_EVENT = "loginHappened";
 
 export default class TabnineAuthenticationProvider
@@ -27,7 +27,7 @@ export default class TabnineAuthenticationProvider
 
   private initializedDisposable: Disposable | undefined;
 
-  private lastState: Promise<State | undefined | null> | undefined;
+  private lastState: State | undefined | null;
 
   private onDidLogin = new Emitter();
 
@@ -44,12 +44,14 @@ export default class TabnineAuthenticationProvider
     );
   }
 
-  async getSessions(): Promise<readonly AuthenticationSession[]> {
-    const state = await this.lastState;
+  getSessions(): Promise<readonly AuthenticationSession[]> {
+    const state = this.lastState;
 
-    return state?.is_logged_in
-      ? [new TabnineSession(state?.user_name, state?.access_token)]
-      : [];
+    return Promise.resolve(
+      state?.is_logged_in
+        ? [new TabnineSession(state?.user_name, state?.access_token)]
+        : []
+    );
   }
 
   async createSession(): Promise<AuthenticationSession> {
@@ -79,36 +81,31 @@ export default class TabnineAuthenticationProvider
     // 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 this.checkForUpdates();
+        void getState().then((state) => {
+          void this.checkForUpdates(state);
+        });
       }
     });
   }
 
   private pollState(): Disposable {
-    let interval: NodeJS.Timeout | undefined;
-    void tabNineProcess.onReady.then(() => {
-      interval = setInterval(() => {
-        void this.checkForUpdates();
-      }, SESSION_POLL_INTERVAL);
-    });
-    return new Disposable(() => {
-      if (interval) {
-        clearInterval(interval);
-      }
+    return BINARY_STATE.onChange((state) => {
+      void this.checkForUpdates(state);
     });
   }
 
-  private async checkForUpdates(): Promise<void> {
+  private async checkForUpdates(
+    state: State | null | undefined
+  ): Promise<void> {
     const added: AuthenticationSession[] = [];
     const removed: AuthenticationSession[] = [];
 
-    const state = getState();
     const { lastState } = this;
 
     this.lastState = state;
 
-    const newState = await this.lastState;
-    const oldState = await lastState;
+    const newState = this.lastState;
+    const oldState = lastState;
 
     if (newState?.is_logged_in) {
       this.onDidLogin.emit(LOGIN_HAPPENED_EVENT, newState);

src/binary/binaryStateSingleton.ts

@@ -0,0 +1,50 @@
+import { Disposable } from "vscode";
+import EventEmitterBasedState from "../state/EventEmitterBasedState";
+import { State } from "./state";
+import { getState, tabNineProcess } from "./requests/requests";
+import { Logger } from "../utils/logger";
+import { BINARY_STATE_POLLING_INTERVAL_MILLISECONDS } from "../globals/consts";
+
+export class BinaryState extends EventEmitterBasedState<State> {
+  private intervalDisposable: Disposable | null = null;
+
+  constructor() {
+    super();
+
+    let interval: NodeJS.Timeout | undefined;
+    void tabNineProcess.onReady.then(() => {
+      interval = setInterval(() => {
+        void this.checkForUpdates();
+      }, BINARY_STATE_POLLING_INTERVAL_MILLISECONDS);
+    });
+
+    this.intervalDisposable = new Disposable(() => {
+      if (interval) {
+        clearInterval(interval);
+      }
+    });
+  }
+
+  private async checkForUpdates() {
+    try {
+      await this.asyncSet(getStateOrNull);
+    } catch (error) {
+      Logger.warn("Failed to refetch state", error);
+    }
+  }
+
+  dispose(): void {
+    super.dispose();
+
+    if (this.intervalDisposable) {
+      this.intervalDisposable.dispose();
+    }
+  }
+}
+
+async function getStateOrNull() {
+  return (await getState()) || null;
+}
+
+const BINARY_STATE = new BinaryState();
+export default BINARY_STATE;

src/enterprise/extension.ts

@@ -41,6 +41,10 @@ import confirmReload from "./update/confirmReload";
 import SignInUsingCustomTokenCommand from "../authentication/loginWithCustomTokenCommand";
 import { SIGN_IN_AUTH_TOKEN_COMMAND } from "../commandsHandler";
 import { WorkspaceUpdater } from "../WorkspaceUpdater";
+import SelfHostedChatEnabledState from "./tabnineChatWidget/SelfHostedChatEnabledState";
+import { emptyStateAuthenticateView } from "../tabnineChatWidget/webviews/emptyStateAuthenticateView";
+import { emptyStateNotPartOfATeamView } from "../tabnineChatWidget/webviews/emptyStateNotPartOfATeamView";
+import BINARY_STATE from "../binary/binaryStateSingleton";
 
 export async function activate(
   context: vscode.ExtensionContext
@@ -49,6 +53,7 @@ export async function activate(
   setTabnineExtensionContext(context);
   context.subscriptions.push(await setEnterpriseContext());
   context.subscriptions.push(new WorkspaceUpdater());
+  context.subscriptions.push(BINARY_STATE);
 
   initReporter(new LogReporter());
   const statusBar = new StatusBar(context);
@@ -86,6 +91,11 @@ export async function activate(
     return;
   }
 
+  context.subscriptions.push(
+    emptyStateAuthenticateView(context),
+    emptyStateNotPartOfATeamView(context)
+  );
+
   const server = serverUrl() as string;
 
   await setBinaryRootPath(context);
@@ -96,7 +106,11 @@ export async function activate(
   }
 
   setBinaryDownloadUrl(server);
-  registerTabnineChatWidgetWebview(context, server);
+  registerTabnineChatWidgetWebview(
+    context,
+    new SelfHostedChatEnabledState(context),
+    server
+  );
 
   await initBinary([
     "--no_bootstrap",

src/enterprise/requests/UserInfo.ts

@@ -2,8 +2,11 @@ import { tabNineProcess } from "../../binary/requests/requests";
 
 export type UserInfo = {
   email: string;
-  team: [];
-  verified: [];
+  team: {
+    id: string;
+    name: string;
+  } | null;
+  verified: boolean;
   isLoggedIn: boolean;
 };
 export default function getUserInfo(): Promise<UserInfo | null | undefined> {

src/enterprise/tabnineChatWidget/SelfHostedChatEnabledState.ts

@@ -0,0 +1,47 @@
+import { ExtensionContext, authentication } from "vscode";
+import ChatEnabledState, {
+  ChatEnabledStateData,
+  ChatStates,
+} from "../../tabnineChatWidget/ChatEnabledState";
+import EventEmitterBasedNonNullState from "../../state/EventEmitterBasedNonNullState";
+import getUserInfo from "../requests/UserInfo";
+
+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();
+      })
+    );
+  }
+
+  async updateState() {
+    await this.asyncSet(fetchChatState);
+  }
+}
+
+async function fetchChatState(): Promise<ChatEnabledStateData | null> {
+  const userInfo = await getUserInfo();
+
+  if (!userInfo) {
+    return null;
+  }
+
+  const isEnabled = userInfo.team !== null;
+
+  if (isEnabled) {
+    return ChatStates.enabled;
+  }
+
+  if (userInfo.isLoggedIn) {
+    return ChatStates.disabled("part_of_a_team_required");
+  }
+
+  return ChatStates.disabled("authnetication_required");
+}

src/extension.ts

@@ -49,13 +49,14 @@ import EventName from "./reports/EventName";
 import registerTabnineChatWidgetWebview from "./tabnineChatWidget/tabnineChatWidgetWebview";
 import { forceRegistrationIfNeeded } from "./registration/forceRegistration";
 import { installationState } from "./events/installationStateChangedEmitter";
-import { statePoller } from "./state/statePoller";
 import { Logger } from "./utils/logger";
 import { callForLogin } from "./authentication/authentication.api";
 import { emptyStateWelcomeView } from "./tabnineChatWidget/webviews/emptyStateChatWelcomeView";
 import { emptyStateAuthenticateView } from "./tabnineChatWidget/webviews/emptyStateAuthenticateView";
 import { activeTextEditorState } from "./activeTextEditorState";
 import { WorkspaceUpdater } from "./WorkspaceUpdater";
+import SaasChatEnabledState from "./tabnineChatWidget/SaasChatEnabledState";
+import BINARY_STATE from "./binary/binaryStateSingleton";
 
 export async function activate(
   context: vscode.ExtensionContext
@@ -67,7 +68,7 @@ export async function activate(
   context.subscriptions.push(handleSelection(context));
   context.subscriptions.push(handleUninstall(() => uponUninstall(context)));
   context.subscriptions.push(installationState);
-  context.subscriptions.push(statePoller);
+  context.subscriptions.push(BINARY_STATE);
   context.subscriptions.push(activeTextEditorState);
   context.subscriptions.push(new WorkspaceUpdater());
   registerCodeReview();
@@ -144,6 +145,7 @@ async function backgroundInit(context: vscode.ExtensionContext) {
 
   registerTabnineChatWidgetWebview(
     context,
+    new SaasChatEnabledState(context),
     context.extensionMode === vscode.ExtensionMode.Test
       ? process.env.CHAT_SERVER_URL
       : undefined

src/globals/consts.ts

@@ -40,7 +40,7 @@ export const BINARY_STATUS_BAR_FIRST_MESSAGE_POLLING_INTERVAL = +(
   process.env.BINARY_NOTIFICATION_POLLING_INTERVAL || 10_000
 ); // 10 seconds
 
-export const BINARY_STATE_POLLING_INTERVAL_MILLISECONDS = 1_000;
+export const BINARY_STATE_POLLING_INTERVAL_MILLISECONDS = 2_000;
 
 export const STATUS_BAR_NOTIFICATION_PERIOD = +(
   process.env.STATUS_BAR_NOTIFICATION_PERIOD || 2 * 60 * 1_000