VSCode client: Add status bar item. (#31)

clients/vscode/package.json

@@ -28,6 +28,10 @@
       {
         "command": "tabby.setServerUrl",
         "title": "Tabby: Set URL of Tabby Server"
+      },
+      {
+        "command": "tabby.openSettings",
+        "title": "Tabby: Open Settings"
       }
     ],
     "configuration": {

clients/vscode/src/Commands.ts

@@ -1,5 +1,5 @@
 import { ConfigurationTarget, workspace, window, commands } from "vscode";
-import { EventHandler } from "./EventHandler";
+import { TabbyClient } from "./TabbyClient";
 
 const target = ConfigurationTarget.Global;
 
@@ -37,14 +37,22 @@ const setServerUrl: Command = {
   },
 };
 
-const eventHandler = new EventHandler();
-const emitEvent: Command = {
-  command: "tabby.emitEvent",
-  callback: (event) => {
-    eventHandler.handle(event);
+const openSettings: Command = {
+  command: "tabby.openSettings",
+  callback: () => {
+    commands.executeCommand("workbench.action.openSettings", "tabby");
   },
 };
 
-export const tabbyCommands = [toogleEnabled, setServerUrl, emitEvent].map((command) =>
+const tabbyClient = TabbyClient.getInstance();
+const emitEvent: Command = {
+  command: "tabby.emitEvent",
+  callback: (event) => {
+    console.debug("Emit Event: ", event);
+    tabbyClient.postEvent(event);
+  },
+};
+
+export const tabbyCommands = [toogleEnabled, setServerUrl, openSettings, emitEvent].map((command) =>
   commands.registerCommand(command.command, command.callback, command.thisArg)
 );

clients/vscode/src/EventHandler.ts

@@ -1,43 +0,0 @@
-import { workspace } from "vscode";
-import axios from "axios";
-
-export enum EventType {
-  InlineCompletionDisplayed,
-  InlineCompletionAccepted,
-}
-
-export interface Event {
-  type: EventType,
-  id?: string,
-  index?: number,
-}
-
-export class EventHandler {
-  private tabbyServerUrl: string = "";
-
-  constructor() {
-    this.updateConfiguration();
-    workspace.onDidChangeConfiguration((event) => {
-      if (event.affectsConfiguration("tabby")) {
-        this.updateConfiguration();
-      }
-    });
-  }
-
-  handle(event: Event) {
-    console.debug("Event: ", event);
-    switch (event.type) {
-      case EventType.InlineCompletionDisplayed:
-        axios.post(`${this.tabbyServerUrl}/v1/completions/${event.id}/choices/${event.index}/view`);
-        break;
-      case EventType.InlineCompletionAccepted:
-        axios.post(`${this.tabbyServerUrl}/v1/completions/${event.id}/choices/${event.index}/select`);
-        break;
-    }
-  }
-
-  private updateConfiguration() {
-    const configuration = workspace.getConfiguration("tabby");
-    this.tabbyServerUrl = configuration.get("serverUrl", "http://127.0.0.1:5000");
-  }
-}

clients/vscode/src/TabbyClient.ts

@@ -0,0 +1,121 @@
+import { workspace } from "vscode";
+import axios from "axios";
+import { sleep } from "./utils";
+import { EventEmitter } from "node:events";
+import { strict as assert } from "node:assert";
+
+const logAxios = false;
+if (logAxios) {
+  axios.interceptors.request.use((request) => {
+    console.debug("Starting Request: ", request);
+    return request;
+  });
+  axios.interceptors.response.use((response) => {
+    console.debug("Response: ", response);
+    return response;
+  });
+}
+
+export interface TabbyCompletion {
+  id?: string;
+  created?: number;
+  choices?: Array<{
+    index: number;
+    text: string;
+  }>;
+}
+
+export enum EventType {
+  InlineCompletionDisplayed = "view",
+  InlineCompletionAccepted = "select",
+}
+
+export interface Event {
+  type: EventType;
+  id?: string;
+  index?: number;
+}
+
+export class TabbyClient extends EventEmitter {
+  private static instance: TabbyClient;
+  static getInstance(): TabbyClient {
+    if (!TabbyClient.instance) {
+      TabbyClient.instance = new TabbyClient();
+    }
+    return TabbyClient.instance;
+  }
+
+  private tabbyServerUrl: string = "";
+  status: "connecting" | "ready" | "disconnected" = "connecting";
+
+  constructor() {
+    super();
+
+    this.updateConfiguration();
+    workspace.onDidChangeConfiguration((event) => {
+      if (event.affectsConfiguration("tabby")) {
+        this.updateConfiguration();
+      }
+    });
+  }
+
+  private updateConfiguration() {
+    const configuration = workspace.getConfiguration("tabby");
+    this.tabbyServerUrl = configuration.get("serverUrl", "http://127.0.0.1:5000");
+    this.ping();
+  }
+
+  private changeStatus(status: "connecting" | "ready" | "disconnected") {
+    if (this.status != status) {
+      this.status = status;
+      this.emit("statusChanged", status);
+    }
+  }
+
+  private async ping(tries: number = 0) {
+    try {
+      const response = await axios.get(`${this.tabbyServerUrl}/`);
+      assert(response.status == 200);
+      this.changeStatus("ready");
+    } catch (e) {
+      if (tries > 5) {
+        this.changeStatus("disconnected");
+        return;
+      }
+      this.changeStatus("connecting");
+      const pingRetryDelay = 1000;
+      await sleep(pingRetryDelay);
+      this.ping(tries + 1);
+    }
+  }
+
+  public async getCompletion(prompt: string): Promise<TabbyCompletion | null> {
+    if (this.status == "disconnected") {
+      this.ping();
+    }
+    try {
+      const response = await axios.post<TabbyCompletion>(`${this.tabbyServerUrl}/v1/completions`, {
+        prompt,
+      });
+      assert(response.status == 200);
+      return response.data;
+    } catch (e) {
+      this.ping();
+      return null;
+    }
+  }
+
+  public async postEvent(event: Event) {
+    if (this.status == "disconnected") {
+      this.ping();
+    }
+    try {
+      const response = await axios.post(
+        `${this.tabbyServerUrl}/v1/completions/${event.id}/choices/${event.index}/${event.type}`
+      );
+      assert(response.status == 200);
+    } catch (e) {
+      this.ping();
+    }
+  }
+}

clients/vscode/src/TabbyCompletionProvider.ts

@@ -10,16 +10,16 @@ import {
   TextDocument,
   workspace,
 } from "vscode";
-import axios, { AxiosResponse } from "axios";
-import { EventType } from "./EventHandler";
+import { TabbyClient, TabbyCompletion, EventType } from "./TabbyClient";
+import { sleep } from "./utils";
 
 export class TabbyCompletionProvider implements InlineCompletionItemProvider {
   private uuid = Date.now();
   private latestTimestamp: number = 0;
 
+  private tabbyClient = TabbyClient.getInstance();
   // User Settings
   private enabled: boolean = true;
-  private tabbyServerUrl: string = "";
 
   constructor() {
     this.updateConfiguration();
@@ -49,7 +49,7 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
     this.latestTimestamp = currentTimestamp;
 
     const suggestionDelay = 150;
-    await this.sleep(suggestionDelay);
+    await sleep(suggestionDelay);
     if (currentTimestamp < this.latestTimestamp) {
       return emptyResponse;
     }
@@ -63,7 +63,7 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
       }
     );
     // Prompt is already nil-checked
-    const response = await this.getCompletions(prompt as String);
+    const completion = await this.tabbyClient.getCompletion(prompt as string);
 
     const hasSuffixParen = this.hasSuffixParen(document, position);
     const replaceRange = hasSuffixParen
@@ -74,7 +74,7 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
           position.character + 1
         )
       : new Range(position, position);
-    const completions = this.toInlineCompletions(response.data, replaceRange);
+    const completions = this.toInlineCompletions(completion, replaceRange);
     console.debug("Result completions: ", completions);
     return Promise.resolve(completions);
   }
@@ -82,7 +82,6 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
   private updateConfiguration() {
     const configuration = workspace.getConfiguration("tabby");
     this.enabled = configuration.get("enabled", true);
-    this.tabbyServerUrl = configuration.get("serverUrl", "http://127.0.0.1:5000");
   }
 
   private getPrompt(document: TextDocument, position: Position): String | undefined {
@@ -96,13 +95,9 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
     return value === undefined || value === null || value.length === 0;
   }
 
-  private sleep(milliseconds: number) {
-    return new Promise((r) => setTimeout(r, milliseconds));
-  }
-
-  private toInlineCompletions(value: any, range: Range): InlineCompletionItem[] {
+  private toInlineCompletions(tabbyCompletion: TabbyCompletion | null, range: Range): InlineCompletionItem[] {
     return (
-      value.choices?.map(
+      tabbyCompletion?.choices?.map(
         (choice: any) =>
           new InlineCompletionItem(choice.text, range, {
             title: "Tabby: Emit Event",
@@ -110,7 +105,7 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
             arguments: [
               {
                 type: EventType.InlineCompletionAccepted,
-                id: value.id,
+                id: tabbyCompletion.id,
                 index: choice.index,
               },
             ],
@@ -119,12 +114,6 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
     );
   }
 
-  private getCompletions(prompt: String): Promise<AxiosResponse<any, any>> {
-    return axios.post(`${this.tabbyServerUrl}/v1/completions`, {
-      prompt,
-    });
-  }
-
   private hasSuffixParen(document: TextDocument, position: Position) {
     const suffix = document.getText(
       new Range(position.line, position.character, position.line, position.character + 1)

clients/vscode/src/TabbyStatusBarItem.ts

@@ -0,0 +1,80 @@
+import { StatusBarAlignment, ThemeColor, window, workspace } from "vscode";
+import { TabbyClient } from "./TabbyClient";
+
+const label = "Tabby";
+const iconLoading = "$(loading~spin)";
+const iconReady = "$(check)";
+const iconDisconnected = "$(plug)";
+const iconDisabled = "$(x)";
+const colorNormal = new ThemeColor('statusBar.foreground');
+const colorWarning = new ThemeColor('statusBarItem.warningForeground');
+const backgroundColorNormal = new ThemeColor('statusBar.background');
+const backgroundColorWarning = new ThemeColor('statusBarItem.warningBackground');
+
+const item = window.createStatusBarItem(StatusBarAlignment.Right);
+export const tabbyStatusBarItem = item;
+
+const client = TabbyClient.getInstance();
+client.on("statusChanged", updateStatusBarItem);
+
+workspace.onDidChangeConfiguration((event) => {
+  if (event.affectsConfiguration("tabby")) {
+    updateStatusBarItem();
+  }
+});
+
+updateStatusBarItem();
+item.show();
+
+function updateStatusBarItem() {
+  const enabled = workspace.getConfiguration("tabby").get("enabled", true);
+  if (!enabled) {
+    toDisabled();
+  } else {
+    const status = client.status;
+    switch (status) {
+      case "connecting":
+        toLoading();
+        break;
+      case "ready":
+        toReady();
+        break;
+      case "disconnected":
+        toDisconnected();
+        break;
+    }
+  }
+}
+
+function toLoading() {
+  item.color = colorNormal;
+  item.backgroundColor = backgroundColorNormal;
+  item.text = `${iconLoading} ${label}`;
+  item.tooltip = "Connecting to Tabby Server...";
+  item.command = undefined;
+}
+
+function toReady() {
+  item.color = colorNormal;
+  item.backgroundColor = backgroundColorNormal;
+  item.text = `${iconReady} ${label}`;
+  item.tooltip = "Tabby is providing code suggestions for you.";
+  item.command = "tabby.toggleEnabled";
+}
+
+function toDisconnected() {
+  item.color = colorWarning;
+  item.backgroundColor = backgroundColorWarning;
+  item.text = `${iconDisconnected} ${label}`;
+  item.tooltip = "Cannot connect to Tabby Server. Click to open settings.";
+  item.command = "tabby.openSettings";
+}
+
+function toDisabled() {
+  item.color = colorWarning;
+  item.backgroundColor = backgroundColorWarning;
+  item.text = `${iconDisabled} ${label}`;
+  item.tooltip = "Tabby is disabled. Click to enable.";
+  item.command = "tabby.toggleEnabled";
+}
+

clients/vscode/src/extension.ts

@@ -3,6 +3,7 @@
 import { ExtensionContext, languages } from "vscode";
 import { tabbyCommands } from "./Commands";
 import { TabbyCompletionProvider } from "./TabbyCompletionProvider";
+import { tabbyStatusBarItem } from "./TabbyStatusBarItem";
 
 // this method is called when your extension is activated
 // your extension is activated the very first time the command is executed
@@ -13,6 +14,7 @@ export function activate(context: ExtensionContext) {
       { pattern: "**" },
       new TabbyCompletionProvider()
     ),
+    tabbyStatusBarItem,
     ...tabbyCommands
   );
 }

clients/vscode/src/utils.ts

@@ -0,0 +1,3 @@
+export function sleep(milliseconds: number) {
+  return new Promise((r) => setTimeout(r, milliseconds));
+}