refactor: improve client auth flow (#269)

* refactor: improve client auth flow

* fix: agent waiting for auth token should update status before resolve.

clients/tabby-agent/dist/index.d.ts

@@ -69,10 +69,21 @@ interface AgentFunction {
     getConfig(): AgentConfig;
     getStatus(): AgentStatus;
     /**
-     * @returns string auth url if AgentStatus is `unauthorized`, null otherwise
+     * @returns the auth url for redirecting, and the code for next step `waitingForAuth`, only return value when
+     *          `AgentStatus` is `unauthorized`, return null otherwise
      * @throws Error if agent is not initialized
      */
-    startAuth(): CancelablePromise<string | null>;
+    requestAuthUrl(): CancelablePromise<{
+        authUrl: string;
+        code: string;
+    } | null>;
+    /**
+     * Wait for auth token to be ready after redirecting user to auth url,
+     * returns nothing, but `AgentStatus` will change to `ready` if resolved successfully
+     * @param code from `requestAuthUrl`
+     * @throws Error if agent is not initialized
+     */
+    waitForAuthToken(code: string): CancelablePromise<any>;
     /**
      * @param request
      * @returns
@@ -94,7 +105,11 @@ type ConfigUpdatedEvent = {
     event: "configUpdated";
     config: AgentConfig;
 };
-type AgentEvent = StatusChangedEvent | ConfigUpdatedEvent;
+type AuthRequiredEvent = {
+    event: "authRequired";
+    server: AgentConfig["server"];
+};
+type AgentEvent = StatusChangedEvent | ConfigUpdatedEvent | AuthRequiredEvent;
 declare const agentEventNames: AgentEvent["event"][];
 interface AgentEventEmitter {
     on<T extends AgentEvent>(eventName: T["event"], callback: (event: T) => void): this;
@@ -145,7 +160,11 @@ declare class TabbyAgent extends EventEmitter implements Agent {
     updateConfig(config: Partial<AgentConfig>): Promise<boolean>;
     getConfig(): AgentConfig;
     getStatus(): AgentStatus;
-    startAuth(): CancelablePromise<string | null>;
+    requestAuthUrl(): CancelablePromise<{
+        authUrl: string;
+        code: string;
+    } | null>;
+    waitForAuthToken(code: string): CancelablePromise<any>;
     getCompletions(request: CompletionRequest): CancelablePromise<CompletionResponse>;
     postEvent(request: LogEventRequest): CancelablePromise<boolean>;
 }

clients/tabby-agent/dist/index.js

@@ -627,14 +627,13 @@ var _Auth = class extends import_events.EventEmitter {
     super();
     this.logger = rootLogger.child({ component: "Auth" });
     this.dataStore = null;
-    this.pollingTokenTimer = null;
-    this.stopPollingTokenTimer = null;
     this.refreshTokenTimer = null;
     this.authApi = null;
     this.jwt = null;
     this.endpoint = options.endpoint;
     this.dataStore = options.dataStore || dataStore;
     this.authApi = new CloudApi();
+    this.scheduleRefreshToken();
   }
   static async create(options) {
     const auth = new _Auth(options);
@@ -665,7 +664,6 @@ var _Auth = class extends import_events.EventEmitter {
         } else {
           this.jwt = jwt;
         }
-        this.scheduleRefreshToken();
       }
     } catch (error) {
       this.logger.debug({ error }, "Error when loading auth");
@@ -695,33 +693,63 @@ var _Auth = class extends import_events.EventEmitter {
       this.jwt = null;
       await this.save();
     }
-    if (this.refreshTokenTimer) {
-      clearTimeout(this.refreshTokenTimer);
-      this.refreshTokenTimer = null;
-    }
-    if (this.pollingTokenTimer) {
-      clearInterval(this.pollingTokenTimer);
-      this.pollingTokenTimer = null;
-    }
-    if (this.stopPollingTokenTimer) {
-      clearTimeout(this.stopPollingTokenTimer);
-      this.stopPollingTokenTimer = null;
-    }
   }
-  async requestToken() {
+  requestAuthUrl() {
-    try {
+    return new CancelablePromise(async (resolve2, reject, onCancel) => {
-      await this.reset();
+      let apiRequest;
-      this.logger.debug("Start to request device token");
+      onCancel(() => {
-      const deviceToken = await this.authApi.api.deviceToken({ auth_url: this.endpoint });
+        apiRequest?.cancel();
-      this.logger.debug({ deviceToken }, "Request device token response");
+      });
-      const authUrl = new URL(_Auth.authPageUrl);
+      try {
-      authUrl.searchParams.append("code", deviceToken.data.code);
+        await this.reset();
-      this.schedulePollingToken(deviceToken.data.code);
+        if (onCancel.isCancelled)
-      return authUrl.toString();
+          return;
-    } catch (error) {
+        this.logger.debug("Start to request device token");
-      this.logger.error({ error }, "Error when requesting token");
+        apiRequest = this.authApi.api.deviceToken({ auth_url: this.endpoint });
-      throw error;
+        const deviceToken = await apiRequest;
-    }
+        this.logger.debug({ deviceToken }, "Request device token response");
+        const authUrl = new URL(_Auth.authPageUrl);
+        authUrl.searchParams.append("code", deviceToken.data.code);
+        resolve2({ authUrl: authUrl.toString(), code: deviceToken.data.code });
+      } catch (error) {
+        this.logger.error({ error }, "Error when requesting token");
+        reject(error);
+      }
+    });
+  }
+  pollingToken(code) {
+    return new CancelablePromise((resolve2, reject, onCancel) => {
+      let apiRequest;
+      const timer = setInterval(async () => {
+        try {
+          apiRequest = this.authApi.api.deviceTokenAccept({ code });
+          const response = await apiRequest;
+          this.logger.debug({ response }, "Poll jwt response");
+          this.jwt = {
+            token: response.data.jwt,
+            payload: (0, import_jwt_decode.default)(response.data.jwt)
+          };
+          super.emit("updated", this.jwt);
+          await this.save();
+          clearInterval(timer);
+          resolve2(true);
+        } catch (error) {
+          if (error instanceof ApiError && [400, 401, 403, 405].indexOf(error.status) !== -1) {
+            this.logger.debug({ error }, "Expected error when polling jwt");
+          } else {
+            this.logger.error({ error }, "Error when polling jwt");
+          }
+        }
+      }, _Auth.tokenStrategy.polling.interval);
+      setTimeout(() => {
+        clearInterval(timer);
+        reject(new Error("Timeout when polling token"));
+      }, _Auth.tokenStrategy.polling.timeout);
+      onCancel(() => {
+        apiRequest?.cancel();
+        clearInterval(timer);
+      });
+    });
   }
   async refreshToken(jwt, options = { maxTry: 1, retryDelay: 1e3 }, retry = 0) {
     try {
@@ -746,50 +774,16 @@ var _Auth = class extends import_events.EventEmitter {
       throw { ...error, retry };
     }
   }
-  async schedulePollingToken(code) {
-    this.pollingTokenTimer = setInterval(async () => {
-      try {
-        const response = await this.authApi.api.deviceTokenAccept({ code });
-        this.logger.debug({ response }, "Poll jwt response");
-        this.jwt = {
-          token: response.data.jwt,
-          payload: (0, import_jwt_decode.default)(response.data.jwt)
-        };
-        await this.save();
-        this.scheduleRefreshToken();
-        super.emit("updated", this.jwt);
-        clearInterval(this.pollingTokenTimer);
-        this.pollingTokenTimer = null;
-      } catch (error) {
-        if (error instanceof ApiError && [400, 401, 403, 405].indexOf(error.status) !== -1) {
-          this.logger.debug({ error }, "Expected error when polling jwt");
-        } else {
-          this.logger.error({ error }, "Error when polling jwt");
-        }
-      }
-    }, _Auth.tokenStrategy.polling.interval);
-    this.stopPollingTokenTimer = setTimeout(() => {
-      if (this.pollingTokenTimer) {
-        clearInterval(this.pollingTokenTimer);
-        this.pollingTokenTimer = null;
-      }
-    }, _Auth.tokenStrategy.polling.timeout);
-  }
   scheduleRefreshToken() {
-    if (this.refreshTokenTimer) {
-      clearTimeout(this.refreshTokenTimer);
-      this.refreshTokenTimer = null;
-    }
-    if (!this.jwt) {
-      return null;
-    }
     this.refreshTokenTimer = setInterval(async () => {
+      if (!this.jwt) {
+        return null;
+      }
       if (this.jwt.payload.exp * 1e3 - Date.now() < _Auth.tokenStrategy.refresh.beforeExpire) {
         try {
           this.jwt = await this.refreshToken(this.jwt, _Auth.tokenStrategy.refresh.whenScheduled);
-          await this.save();
-          this.scheduleRefreshToken();
           super.emit("updated", this.jwt);
+          await this.save();
         } catch (error) {
           this.logger.error({ error }, "Error when refreshing jwt");
         }
@@ -1211,6 +1205,10 @@ var _TabbyAgent = class extends import_events2.EventEmitter {
       this.config = (0, import_deepmerge.default)(this.config, options.config);
     }
     await this.applyConfig();
+    if (this.status === "unauthorized") {
+      const event = { event: "authRequired", server: this.config.server };
+      super.emit("authRequired", event);
+    }
     await this.anonymousUsageLogger.event("AgentInitialized", {
       client: options.client
     });
@@ -1220,11 +1218,16 @@ var _TabbyAgent = class extends import_events2.EventEmitter {
   async updateConfig(config) {
     const mergedConfig = (0, import_deepmerge.default)(this.config, config);
     if (!(0, import_deep_equal.default)(this.config, mergedConfig)) {
+      const serverUpdated = !(0, import_deep_equal.default)(this.config.server, mergedConfig.server);
       this.config = mergedConfig;
       await this.applyConfig();
       const event = { event: "configUpdated", config: this.config };
       this.logger.debug({ event }, "Config updated");
       super.emit("configUpdated", event);
+      if (serverUpdated && this.status === "unauthorized") {
+        const event2 = { event: "authRequired", server: this.config.server };
+        super.emit("authRequired", event2);
+      }
     }
     return true;
   }
@@ -1234,27 +1237,46 @@ var _TabbyAgent = class extends import_events2.EventEmitter {
   getStatus() {
     return this.status;
   }
-  startAuth() {
+  requestAuthUrl() {
     if (this.status === "notInitialized") {
-      throw new Error("Agent is not initialized");
+      return cancelable(Promise.reject("Agent is not initialized"), () => {
+      });
     }
+    return new CancelablePromise(async (resolve2, reject, onCancel) => {
+      let request2;
+      onCancel(() => {
+        request2?.cancel();
+      });
+      await this.healthCheck();
+      if (onCancel.isCancelled)
+        return;
+      if (this.status === "unauthorized") {
+        request2 = this.auth.requestAuthUrl();
+        resolve2(request2);
+      } else {
+      }
+      resolve2(null);
+    });
+  }
+  waitForAuthToken(code) {
+    if (this.status === "notInitialized") {
+      return cancelable(Promise.reject("Agent is not initialized"), () => {
+      });
+    }
+    const polling = this.auth.pollingToken(code);
     return cancelable(
-      this.healthCheck().then(() => {
+      polling.then(() => {
-        if (this.status === "unauthorized") {
+        return this.setupApi();
-          return this.auth.requestToken();
-        }
-        return null;
       }),
       () => {
-        if (this.status === "unauthorized") {
+        polling.cancel();
-          this.auth.reset();
-        }
       }
     );
   }
   getCompletions(request2) {
     if (this.status === "notInitialized") {
-      throw new Error("Agent is not initialized");
+      return cancelable(Promise.reject("Agent is not initialized"), () => {
+      });
     }
     if (this.completionCache.has(request2)) {
       this.logger.debug({ request: request2 }, "Completion cache hit");
@@ -1291,7 +1313,8 @@ var _TabbyAgent = class extends import_events2.EventEmitter {
   }
   postEvent(request2) {
     if (this.status === "notInitialized") {
-      throw new Error("Agent is not initialized");
+      return cancelable(Promise.reject("Agent is not initialized"), () => {
+      });
     }
     return this.callApi(this.api.v1.event, request2);
   }
@@ -1300,7 +1323,7 @@ var TabbyAgent = _TabbyAgent;
 TabbyAgent.tryConnectInterval = 1e3 * 30;
 
 // src/Agent.ts
-var agentEventNames = ["statusChanged", "configUpdated"];
+var agentEventNames = ["statusChanged", "configUpdated", "authRequired"];
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
   CancelablePromise,

clients/tabby-agent/dist/index.mjs

@@ -13353,14 +13353,13 @@ var _Auth = class extends EventEmitter {
     super();
     this.logger = rootLogger.child({ component: "Auth" });
     this.dataStore = null;
-    this.pollingTokenTimer = null;
-    this.stopPollingTokenTimer = null;
     this.refreshTokenTimer = null;
     this.authApi = null;
     this.jwt = null;
     this.endpoint = options.endpoint;
     this.dataStore = options.dataStore || dataStore;
     this.authApi = new CloudApi();
+    this.scheduleRefreshToken();
   }
   static async create(options) {
     const auth = new _Auth(options);
@@ -13391,7 +13390,6 @@ var _Auth = class extends EventEmitter {
         } else {
           this.jwt = jwt;
         }
-        this.scheduleRefreshToken();
       }
     } catch (error) {
       this.logger.debug({ error }, "Error when loading auth");
@@ -13421,33 +13419,63 @@ var _Auth = class extends EventEmitter {
       this.jwt = null;
       await this.save();
     }
-    if (this.refreshTokenTimer) {
-      clearTimeout(this.refreshTokenTimer);
-      this.refreshTokenTimer = null;
-    }
-    if (this.pollingTokenTimer) {
-      clearInterval(this.pollingTokenTimer);
-      this.pollingTokenTimer = null;
-    }
-    if (this.stopPollingTokenTimer) {
-      clearTimeout(this.stopPollingTokenTimer);
-      this.stopPollingTokenTimer = null;
-    }
   }
-  async requestToken() {
+  requestAuthUrl() {
-    try {
+    return new CancelablePromise(async (resolve4, reject, onCancel) => {
-      await this.reset();
+      let apiRequest;
-      this.logger.debug("Start to request device token");
+      onCancel(() => {
-      const deviceToken = await this.authApi.api.deviceToken({ auth_url: this.endpoint });
+        apiRequest?.cancel();
-      this.logger.debug({ deviceToken }, "Request device token response");
+      });
-      const authUrl = new URL(_Auth.authPageUrl);
+      try {
-      authUrl.searchParams.append("code", deviceToken.data.code);
+        await this.reset();
-      this.schedulePollingToken(deviceToken.data.code);
+        if (onCancel.isCancelled)
-      return authUrl.toString();
+          return;
-    } catch (error) {
+        this.logger.debug("Start to request device token");
-      this.logger.error({ error }, "Error when requesting token");
+        apiRequest = this.authApi.api.deviceToken({ auth_url: this.endpoint });
-      throw error;
+        const deviceToken = await apiRequest;
-    }
+        this.logger.debug({ deviceToken }, "Request device token response");
+        const authUrl = new URL(_Auth.authPageUrl);
+        authUrl.searchParams.append("code", deviceToken.data.code);
+        resolve4({ authUrl: authUrl.toString(), code: deviceToken.data.code });
+      } catch (error) {
+        this.logger.error({ error }, "Error when requesting token");
+        reject(error);
+      }
+    });
+  }
+  pollingToken(code) {
+    return new CancelablePromise((resolve4, reject, onCancel) => {
+      let apiRequest;
+      const timer = setInterval(async () => {
+        try {
+          apiRequest = this.authApi.api.deviceTokenAccept({ code });
+          const response = await apiRequest;
+          this.logger.debug({ response }, "Poll jwt response");
+          this.jwt = {
+            token: response.data.jwt,
+            payload: jwt_decode_esm_default(response.data.jwt)
+          };
+          super.emit("updated", this.jwt);
+          await this.save();
+          clearInterval(timer);
+          resolve4(true);
+        } catch (error) {
+          if (error instanceof ApiError && [400, 401, 403, 405].indexOf(error.status) !== -1) {
+            this.logger.debug({ error }, "Expected error when polling jwt");
+          } else {
+            this.logger.error({ error }, "Error when polling jwt");
+          }
+        }
+      }, _Auth.tokenStrategy.polling.interval);
+      setTimeout(() => {
+        clearInterval(timer);
+        reject(new Error("Timeout when polling token"));
+      }, _Auth.tokenStrategy.polling.timeout);
+      onCancel(() => {
+        apiRequest?.cancel();
+        clearInterval(timer);
+      });
+    });
   }
   async refreshToken(jwt, options = { maxTry: 1, retryDelay: 1e3 }, retry = 0) {
     try {
@@ -13472,50 +13500,16 @@ var _Auth = class extends EventEmitter {
       throw { ...error, retry };
     }
   }
-  async schedulePollingToken(code) {
-    this.pollingTokenTimer = setInterval(async () => {
-      try {
-        const response = await this.authApi.api.deviceTokenAccept({ code });
-        this.logger.debug({ response }, "Poll jwt response");
-        this.jwt = {
-          token: response.data.jwt,
-          payload: jwt_decode_esm_default(response.data.jwt)
-        };
-        await this.save();
-        this.scheduleRefreshToken();
-        super.emit("updated", this.jwt);
-        clearInterval(this.pollingTokenTimer);
-        this.pollingTokenTimer = null;
-      } catch (error) {
-        if (error instanceof ApiError && [400, 401, 403, 405].indexOf(error.status) !== -1) {
-          this.logger.debug({ error }, "Expected error when polling jwt");
-        } else {
-          this.logger.error({ error }, "Error when polling jwt");
-        }
-      }
-    }, _Auth.tokenStrategy.polling.interval);
-    this.stopPollingTokenTimer = setTimeout(() => {
-      if (this.pollingTokenTimer) {
-        clearInterval(this.pollingTokenTimer);
-        this.pollingTokenTimer = null;
-      }
-    }, _Auth.tokenStrategy.polling.timeout);
-  }
   scheduleRefreshToken() {
-    if (this.refreshTokenTimer) {
-      clearTimeout(this.refreshTokenTimer);
-      this.refreshTokenTimer = null;
-    }
-    if (!this.jwt) {
-      return null;
-    }
     this.refreshTokenTimer = setInterval(async () => {
+      if (!this.jwt) {
+        return null;
+      }
       if (this.jwt.payload.exp * 1e3 - Date.now() < _Auth.tokenStrategy.refresh.beforeExpire) {
         try {
           this.jwt = await this.refreshToken(this.jwt, _Auth.tokenStrategy.refresh.whenScheduled);
-          await this.save();
-          this.scheduleRefreshToken();
           super.emit("updated", this.jwt);
+          await this.save();
         } catch (error) {
           this.logger.error({ error }, "Error when refreshing jwt");
         }
@@ -15291,6 +15285,10 @@ var _TabbyAgent = class extends EventEmitter {
       this.config = (0, import_deepmerge.default)(this.config, options.config);
     }
     await this.applyConfig();
+    if (this.status === "unauthorized") {
+      const event = { event: "authRequired", server: this.config.server };
+      super.emit("authRequired", event);
+    }
     await this.anonymousUsageLogger.event("AgentInitialized", {
       client: options.client
     });
@@ -15300,11 +15298,16 @@ var _TabbyAgent = class extends EventEmitter {
   async updateConfig(config2) {
     const mergedConfig = (0, import_deepmerge.default)(this.config, config2);
     if (!(0, import_deep_equal.default)(this.config, mergedConfig)) {
+      const serverUpdated = !(0, import_deep_equal.default)(this.config.server, mergedConfig.server);
       this.config = mergedConfig;
       await this.applyConfig();
       const event = { event: "configUpdated", config: this.config };
       this.logger.debug({ event }, "Config updated");
       super.emit("configUpdated", event);
+      if (serverUpdated && this.status === "unauthorized") {
+        const event2 = { event: "authRequired", server: this.config.server };
+        super.emit("authRequired", event2);
+      }
     }
     return true;
   }
@@ -15314,27 +15317,45 @@ var _TabbyAgent = class extends EventEmitter {
   getStatus() {
     return this.status;
   }
-  startAuth() {
+  requestAuthUrl() {
     if (this.status === "notInitialized") {
-      throw new Error("Agent is not initialized");
+      return cancelable(Promise.reject("Agent is not initialized"), () => {
+      });
     }
+    return new CancelablePromise(async (resolve4, reject, onCancel) => {
+      let request2;
+      onCancel(() => {
+        request2?.cancel();
+      });
+      await this.healthCheck();
+      if (onCancel.isCancelled)
+        return;
+      if (this.status === "unauthorized") {
+        request2 = this.auth.requestAuthUrl();
+        resolve4(request2);
+      }
+      resolve4(null);
+    });
+  }
+  waitForAuthToken(code) {
+    if (this.status === "notInitialized") {
+      return cancelable(Promise.reject("Agent is not initialized"), () => {
+      });
+    }
+    const polling = this.auth.pollingToken(code);
     return cancelable(
-      this.healthCheck().then(() => {
+      polling.then(() => {
-        if (this.status === "unauthorized") {
+        return this.setupApi();
-          return this.auth.requestToken();
-        }
-        return null;
       }),
       () => {
-        if (this.status === "unauthorized") {
+        polling.cancel();
-          this.auth.reset();
-        }
       }
     );
   }
   getCompletions(request2) {
     if (this.status === "notInitialized") {
-      throw new Error("Agent is not initialized");
+      return cancelable(Promise.reject("Agent is not initialized"), () => {
+      });
     }
     if (this.completionCache.has(request2)) {
       this.logger.debug({ request: request2 }, "Completion cache hit");
@@ -15371,7 +15392,8 @@ var _TabbyAgent = class extends EventEmitter {
   }
   postEvent(request2) {
     if (this.status === "notInitialized") {
-      throw new Error("Agent is not initialized");
+      return cancelable(Promise.reject("Agent is not initialized"), () => {
+      });
     }
     return this.callApi(this.api.v1.event, request2);
   }
@@ -15385,7 +15407,7 @@ init_dirname();
 init_filename();
 init_buffer2();
 init_process2();
-var agentEventNames = ["statusChanged", "configUpdated"];
+var agentEventNames = ["statusChanged", "configUpdated", "authRequired"];
 /*! Bundled license information:
 
 @jspm/core/nodelibs/browser/buffer.js:

clients/tabby-agent/src/Agent.ts

@@ -28,25 +28,34 @@ export type AgentStatus = "notInitialized" | "ready" | "disconnected" | "unautho
 
 export interface AgentFunction {
   initialize(options: Partial<AgentInitOptions>): Promise<boolean>;
-  updateConfig(config: Partial<AgentConfig>):  Promise<boolean>;
+  updateConfig(config: Partial<AgentConfig>): Promise<boolean>;
   getConfig(): AgentConfig;
   getStatus(): AgentStatus;
 
   /**
-   * @returns string auth url if AgentStatus is `unauthorized`, null otherwise
+   * @returns the auth url for redirecting, and the code for next step `waitingForAuth`, only return value when
+   *          `AgentStatus` is `unauthorized`, return null otherwise
    * @throws Error if agent is not initialized
    */
-  startAuth(): CancelablePromise<string | null>;
+  requestAuthUrl(): CancelablePromise<{ authUrl: string; code: string } | null>;
 
   /**
-   * @param request 
+   * Wait for auth token to be ready after redirecting user to auth url,
+   * returns nothing, but `AgentStatus` will change to `ready` if resolved successfully
+   * @param code from `requestAuthUrl`
+   * @throws Error if agent is not initialized
+   */
+  waitForAuthToken(code: string): CancelablePromise<any>;
+
+  /**
+   * @param request
    * @returns
    * @throws Error if agent is not initialized
    */
   getCompletions(request: CompletionRequest): CancelablePromise<CompletionResponse>;
 
   /**
-   * @param event 
+   * @param event
    * @returns
    * @throws Error if agent is not initialized
    */
@@ -61,9 +70,13 @@ export type ConfigUpdatedEvent = {
   event: "configUpdated";
   config: AgentConfig;
 };
+export type AuthRequiredEvent = {
+  event: "authRequired";
+  server: AgentConfig["server"]
+};
 
-export type AgentEvent = StatusChangedEvent | ConfigUpdatedEvent;
+export type AgentEvent = StatusChangedEvent | ConfigUpdatedEvent | AuthRequiredEvent;
-export const agentEventNames: AgentEvent["event"][] = ["statusChanged", "configUpdated"];
+export const agentEventNames: AgentEvent["event"][] = ["statusChanged", "configUpdated", "authRequired"];
 
 export interface AgentEventEmitter {
   on<T extends AgentEvent>(eventName: T["event"], callback: (event: T) => void): this;

clients/tabby-agent/src/Auth.ts

@@ -1,7 +1,7 @@
 import { EventEmitter } from "events";
 import decodeJwt from "jwt-decode";
-import { CloudApi } from "./cloud";
+import { CloudApi, DeviceTokenResponse, DeviceTokenAcceptResponse } from "./cloud";
-import { ApiError } from "./generated";
+import { ApiError, CancelablePromise } from "./generated";
 import { dataStore, DataStore } from "./dataStore";
 import { rootLogger } from "./logger";
 
@@ -39,8 +39,6 @@ export class Auth extends EventEmitter {
   private readonly logger = rootLogger.child({ component: "Auth" });
   readonly endpoint: string;
   readonly dataStore: DataStore | null = null;
-  private pollingTokenTimer: ReturnType<typeof setInterval> | null = null;
-  private stopPollingTokenTimer: ReturnType<typeof setTimeout> | null = null;
   private refreshTokenTimer: ReturnType<typeof setInterval> | null = null;
   private authApi: CloudApi | null = null;
   private jwt: JWT | null = null;
@@ -56,6 +54,7 @@ export class Auth extends EventEmitter {
     this.endpoint = options.endpoint;
     this.dataStore = options.dataStore || dataStore;
     this.authApi = new CloudApi();
+    this.scheduleRefreshToken();
   }
 
   get token(): string | null {
@@ -84,7 +83,6 @@ export class Auth extends EventEmitter {
         } else {
           this.jwt = jwt;
         }
-        this.scheduleRefreshToken();
       }
     } catch (error: any) {
       this.logger.debug({ error }, "Error when loading auth");
@@ -113,34 +111,65 @@ export class Auth extends EventEmitter {
       this.jwt = null;
       await this.save();
     }
-    if (this.refreshTokenTimer) {
-      clearTimeout(this.refreshTokenTimer);
-      this.refreshTokenTimer = null;
-    }
-    if (this.pollingTokenTimer) {
-      clearInterval(this.pollingTokenTimer);
-      this.pollingTokenTimer = null;
-    }
-    if (this.stopPollingTokenTimer) {
-      clearTimeout(this.stopPollingTokenTimer);
-      this.stopPollingTokenTimer = null;
-    }
   }
 
-  async requestToken(): Promise<string> {
+  requestAuthUrl(): CancelablePromise<{ authUrl: string; code: string }> {
-    try {
+    return new CancelablePromise(async (resolve, reject, onCancel) => {
-      await this.reset();
+      let apiRequest: CancelablePromise<DeviceTokenResponse>;
-      this.logger.debug("Start to request device token");
+      onCancel(() => {
-      const deviceToken = await this.authApi.api.deviceToken({ auth_url: this.endpoint });
+        apiRequest?.cancel();
-      this.logger.debug({ deviceToken }, "Request device token response");
+      });
-      const authUrl = new URL(Auth.authPageUrl);
+      try {
-      authUrl.searchParams.append("code", deviceToken.data.code);
+        await this.reset();
-      this.schedulePollingToken(deviceToken.data.code);
+        if (onCancel.isCancelled) return;
-      return authUrl.toString();
+        this.logger.debug("Start to request device token");
-    } catch (error) {
+        apiRequest = this.authApi.api.deviceToken({ auth_url: this.endpoint });
-      this.logger.error({ error }, "Error when requesting token");
+        const deviceToken = await apiRequest;
-      throw error;
+        this.logger.debug({ deviceToken }, "Request device token response");
-    }
+        const authUrl = new URL(Auth.authPageUrl);
+        authUrl.searchParams.append("code", deviceToken.data.code);
+        resolve({ authUrl: authUrl.toString(), code: deviceToken.data.code });
+      } catch (error) {
+        this.logger.error({ error }, "Error when requesting token");
+        reject(error);
+      }
+    });
+  }
+
+  pollingToken(code: string): CancelablePromise<boolean> {
+    return new CancelablePromise((resolve, reject, onCancel) => {
+      let apiRequest: CancelablePromise<DeviceTokenAcceptResponse>;
+      const timer = setInterval(async () => {
+        try {
+          apiRequest = this.authApi.api.deviceTokenAccept({ code });
+          const response = await apiRequest;
+          this.logger.debug({ response }, "Poll jwt response");
+          this.jwt = {
+            token: response.data.jwt,
+            payload: decodeJwt(response.data.jwt),
+          };
+          super.emit("updated", this.jwt);
+          await this.save();
+          clearInterval(timer);
+          resolve(true);
+        } catch (error) {
+          if (error instanceof ApiError && [400, 401, 403, 405].indexOf(error.status) !== -1) {
+            this.logger.debug({ error }, "Expected error when polling jwt");
+          } else {
+            // unknown error but still keep polling
+            this.logger.error({ error }, "Error when polling jwt");
+          }
+        }
+      }, Auth.tokenStrategy.polling.interval);
+      setTimeout(() => {
+        clearInterval(timer);
+        reject(new Error("Timeout when polling token"));
+      }, Auth.tokenStrategy.polling.timeout);
+      onCancel(() => {
+        apiRequest?.cancel();
+        clearInterval(timer);
+      });
+    });
   }
 
   private async refreshToken(jwt: JWT, options = { maxTry: 1, retryDelay: 1000 }, retry = 0): Promise<JWT> {
@@ -168,53 +197,16 @@ export class Auth extends EventEmitter {
     }
   }
 
-  private async schedulePollingToken(code: string) {
-    this.pollingTokenTimer = setInterval(async () => {
-      try {
-        const response = await this.authApi.api.deviceTokenAccept({ code });
-        this.logger.debug({ response }, "Poll jwt response");
-        this.jwt = {
-          token: response.data.jwt,
-          payload: decodeJwt(response.data.jwt),
-        };
-        await this.save();
-        this.scheduleRefreshToken();
-        super.emit("updated", this.jwt);
-        clearInterval(this.pollingTokenTimer);
-        this.pollingTokenTimer = null;
-      } catch (error) {
-        if (error instanceof ApiError && [400, 401, 403, 405].indexOf(error.status) !== -1) {
-          this.logger.debug({ error }, "Expected error when polling jwt");
-        } else {
-          // unknown error but still keep polling
-          this.logger.error({ error }, "Error when polling jwt");
-        }
-      }
-    }, Auth.tokenStrategy.polling.interval);
-    this.stopPollingTokenTimer = setTimeout(() => {
-      if (this.pollingTokenTimer) {
-        clearInterval(this.pollingTokenTimer);
-        this.pollingTokenTimer = null;
-      }
-    }, Auth.tokenStrategy.polling.timeout);
-  }
-
   private scheduleRefreshToken() {
-    if (this.refreshTokenTimer) {
-      clearTimeout(this.refreshTokenTimer);
-      this.refreshTokenTimer = null;
-    }
-    if (!this.jwt) {
-      return null;
-    }
-
     this.refreshTokenTimer = setInterval(async () => {
+      if (!this.jwt) {
+        return null;
+      }
       if (this.jwt.payload.exp * 1000 - Date.now() < Auth.tokenStrategy.refresh.beforeExpire) {
         try {
           this.jwt = await this.refreshToken(this.jwt, Auth.tokenStrategy.refresh.whenScheduled);
-          await this.save();
-          this.scheduleRefreshToken();
           super.emit("updated", this.jwt);
+          await this.save();
         } catch (error) {
           this.logger.error({ error }, "Error when refreshing jwt");
         }

clients/tabby-agent/src/TabbyAgent.ts

@@ -148,6 +148,10 @@ export class TabbyAgent extends EventEmitter implements Agent {
       this.config = deepMerge(this.config, options.config);
     }
     await this.applyConfig();
+    if (this.status === "unauthorized") {
+      const event: AgentEvent = { event: "authRequired", server: this.config.server };
+      super.emit("authRequired", event);
+    }
     await this.anonymousUsageLogger.event("AgentInitialized", {
       client: options.client,
     });
@@ -158,11 +162,16 @@ export class TabbyAgent extends EventEmitter implements Agent {
   public async updateConfig(config: Partial<AgentConfig>): Promise<boolean> {
     const mergedConfig = deepMerge(this.config, config);
     if (!deepEqual(this.config, mergedConfig)) {
+      const serverUpdated = !deepEqual(this.config.server, mergedConfig.server);
       this.config = mergedConfig;
       await this.applyConfig();
       const event: AgentEvent = { event: "configUpdated", config: this.config };
       this.logger.debug({ event }, "Config updated");
       super.emit("configUpdated", event);
+      if (serverUpdated && this.status === "unauthorized") {
+        const event: AgentEvent = { event: "authRequired", server: this.config.server };
+        super.emit("authRequired", event);
+      }
     }
     return true;
   }
@@ -175,28 +184,44 @@ export class TabbyAgent extends EventEmitter implements Agent {
     return this.status;
   }
 
-  public startAuth(): CancelablePromise<string | null> {
+  public requestAuthUrl(): CancelablePromise<{ authUrl: string; code: string } | null> {
     if (this.status === "notInitialized") {
-      throw new Error("Agent is not initialized");
+      return cancelable(Promise.reject("Agent is not initialized"), () => {});
     }
+    return new CancelablePromise(async (resolve, reject, onCancel) => {
+      let request: CancelablePromise<{ authUrl: string; code: string }>;
+      onCancel(() => {
+        request?.cancel();
+      });
+      await this.healthCheck();
+      if (onCancel.isCancelled) return;
+      if (this.status === "unauthorized") {
+        request = this.auth.requestAuthUrl();
+        resolve(request);
+      } else {
+      }
+      resolve(null);
+    });
+  }
+
+  public waitForAuthToken(code: string): CancelablePromise<any> {
+    if (this.status === "notInitialized") {
+      return cancelable(Promise.reject("Agent is not initialized"), () => {});
+    }
+    const polling = this.auth.pollingToken(code);
     return cancelable(
-      this.healthCheck().then(() => {
+      polling.then(() => {
-        if (this.status === "unauthorized") {
+        return this.setupApi();
-          return this.auth.requestToken();
-        }
-        return null;
       }),
       () => {
-        if (this.status === "unauthorized") {
+        polling.cancel();
-          this.auth.reset();
-        }
       }
     );
   }
 
   public getCompletions(request: CompletionRequest): CancelablePromise<CompletionResponse> {
     if (this.status === "notInitialized") {
-      throw new Error("Agent is not initialized");
+      return cancelable(Promise.reject("Agent is not initialized"), () => {});
     }
     if (this.completionCache.has(request)) {
       this.logger.debug({ request }, "Completion cache hit");
@@ -236,7 +261,7 @@ export class TabbyAgent extends EventEmitter implements Agent {
 
   public postEvent(request: LogEventRequest): CancelablePromise<boolean> {
     if (this.status === "notInitialized") {
-      throw new Error("Agent is not initialized");
+      return cancelable(Promise.reject("Agent is not initialized"), () => {});
     }
     return this.callApi(this.api.v1.event, request);
   }

clients/vscode/src/commands.ts

@@ -1,6 +1,7 @@
 import {
   ConfigurationTarget,
   InputBoxValidationSeverity,
+  ProgressLocation,
   QuickPickItem,
   QuickPickItemKind,
   Uri,
@@ -11,6 +12,7 @@ import {
 } from "vscode";
 import { strict as assert } from "assert";
 import { Duration } from "@sapphire/duration";
+import { CancelablePromise } from "tabby-agent";
 import { agent } from "./agent";
 import { notifications } from "./notifications";
 
@@ -140,23 +142,46 @@ const emitEvent: Command = {
 
 const openAuthPage: Command = {
   command: "tabby.openAuthPage",
-  callback: (callbacks?: { onOpenAuthPage?: () => void }) => {
+  callback: (callbacks?: { onAuthStart?: () => void; onAuthEnd?: () => void }) => {
-    agent()
+    window.withProgress(
-      .startAuth()
+      {
-      .then((authUrl) => {
+        location: ProgressLocation.Notification,
-        if (authUrl) {
+        title: "Tabby Server Authorization",
-          callbacks?.onOpenAuthPage?.();
+        cancellable: true,
-          env.openExternal(Uri.parse(authUrl));
+      },
-        } else if (agent().getStatus() === "ready") {
+      async (progress, token) => {
-          notifications.showInformationWhenStartAuthButAlreadyAuthorized();
+        let requestAuthUrl: CancelablePromise<{ authUrl: string; code: string } | null>;
-        } else {
+        let waitForAuthToken: CancelablePromise<any>;
-          notifications.showInformationWhenStartAuthFailed();
+        token.onCancellationRequested(() => {
+          requestAuthUrl?.cancel();
+          waitForAuthToken?.cancel();
+        });
+        try {
+          callbacks?.onAuthStart?.();
+          progress.report({ message: "Generating authorization url..." });
+          requestAuthUrl = agent().requestAuthUrl();
+          let authUrl = await requestAuthUrl;
+          if (authUrl) {
+            env.openExternal(Uri.parse(authUrl.authUrl));
+            progress.report({ message: "Waiting for authorization from browser..." });
+            waitForAuthToken = agent().waitForAuthToken(authUrl.code);
+            await waitForAuthToken;
+            assert(agent().getStatus() === "ready");
+            notifications.showInformationAuthSuccess();
+          } else if (agent().getStatus() === "ready") {
+            notifications.showInformationWhenStartAuthButAlreadyAuthorized();
+          } else {
+            notifications.showInformationWhenAuthFailed();
+          }
+        } catch (error: any) {
+          if (error.isCancelled) return;
+          console.debug("Error auth", { error });
+          notifications.showInformationWhenAuthFailed();
+        } finally {
+          callbacks?.onAuthEnd?.();
         }
-      })
+      }
-      .catch((error) => {
+    );
-        console.debug("Error to start auth", { error });
-        notifications.showInformationWhenStartAuthFailed();
-      });
   },
 };
 

clients/vscode/src/notifications.ts

@@ -52,9 +52,9 @@ function showInformationWhenDisconnected() {
     });
 }
 
-function showInformationStartAuth(callbacks?: { onOpenAuthPage?: () => void }) {
+function showInformationStartAuth(callbacks?: { onAuthStart?: () => void; onAuthEnd?: () => void }) {
   window
-    .showInformationMessage(
+    .showWarningMessage(
       "Tabby Server requires authorization. Continue to open authorization page in your browser.",
       "Continue",
       "Settings"
@@ -78,8 +78,8 @@ function showInformationWhenStartAuthButAlreadyAuthorized() {
   window.showInformationMessage("You are already authorized now.");
 }
 
-function showInformationWhenStartAuthFailed() {
+function showInformationWhenAuthFailed() {
-  window.showInformationMessage("Cannot connect to server. Please check settings.", "Settings").then((selection) => {
+  window.showWarningMessage("Cannot connect to server. Please check settings.", "Settings").then((selection) => {
     switch (selection) {
       case "Settings":
         commands.executeCommand("tabby.openSettings");
@@ -96,5 +96,5 @@ export const notifications = {
   showInformationStartAuth,
   showInformationAuthSuccess,
   showInformationWhenStartAuthButAlreadyAuthorized,
-  showInformationWhenStartAuthFailed,
+  showInformationWhenAuthFailed,
 };

clients/vscode/src/statusBarItem.ts

@@ -36,24 +36,14 @@ const fsm = createMachine({
         ready: "ready",
         disconnected: "disconnected",
         disabled: "disabled",
-        openAuthPage: "unauthorizedAndAuthPageOpen",
+        authStart: "unauthorizedAndAuthInProgress",
-      },
-      entry: () => {
-        toUnauthorized();
-        notifications.showInformationStartAuth({
-          onOpenAuthPage: () => {
-            fsmService.send("openAuthPage");
-          },
-        });
       },
+      entry: () => toUnauthorized(),
     },
-    unauthorizedAndAuthPageOpen: {
+    unauthorizedAndAuthInProgress: {
-      on: { ready: "ready", disconnected: "disconnected", disabled: "disabled" },
+      // if auth succeeds, we will get `ready` before `authEnd` event
-      exit: (_, event) => {
+      on: { ready: "ready", disconnected: "disconnected", disabled: "disabled", authEnd: "unauthorized" },
-        if (event.type === "ready") {
+      entry: () => toUnauthorizedAndAuthInProgress(),
-          notifications.showInformationAuthSuccess();
-        }
-      },
     },
     disabled: {
       on: { loading: "loading", ready: "ready", disconnected: "disconnected", unauthorized: "unauthorized" },
@@ -95,6 +85,14 @@ function toUnauthorized() {
   item.command = { title: "", command: "tabby.statusBarItemClicked", arguments: ["unauthorized"] };
 }
 
+function toUnauthorizedAndAuthInProgress() {
+  item.color = colorWarning;
+  item.backgroundColor = backgroundColorWarning;
+  item.text = `${iconUnauthorized} ${label}`;
+  item.tooltip = "Waiting for authorization.";
+  item.command = undefined;
+}
+
 function toDisabled() {
   item.color = colorWarning;
   item.backgroundColor = backgroundColorWarning;
@@ -133,6 +131,17 @@ export const tabbyStatusBarItem = () => {
   });
   agent().on("statusChanged", updateStatusBarItem);
 
+  agent().on("authRequired", () => {
+    notifications.showInformationStartAuth({
+      onAuthStart: () => {
+        fsmService.send("authStart");
+      },
+      onAuthEnd: () => {
+        fsmService.send("authEnd");
+      },
+    });
+  });
+
   item.show();
   return item;
 };