Skip to content
Draft
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 25 additions & 9 deletions webapp/_webapp/src/libs/apiclient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import { EventEmitter } from "events";
import { ErrorCode, ErrorSchema } from "../pkg/gen/apiclient/shared/v1/shared_pb";
import { errorToast } from "./toasts";
import { storage } from "./storage";
import { useAuthStore } from "../stores/auth-store";

// Exhaustive type check helper - will cause compile error if a case is not handled
const assertNever = (x: never): never => {
Expand All @@ -29,6 +30,7 @@ class ApiClient {
private axiosInstance: AxiosInstance;
private refreshToken: string | null;
private onTokenRefreshedEventEmitter: EventEmitter;
private refreshPromise: Promise<void> | null = null;

constructor(baseURL: string, apiVersion: ApiVersion) {
this.axiosInstance = axios.create({
Expand Down Expand Up @@ -64,6 +66,8 @@ class ApiClient {
}

setTokens(token: string, refreshToken: string): void {
useAuthStore.getState().setToken(token);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

May i check if im understanding the situation right. This fix is for 401 error?

In particular, when we refresh due to 401 error from sending a request with expired tokens, we get the new generated tokens but this was never updated to auth-store (browser localStorage). And so, on the next session, e.g. page reload, auth-store loads stale tokens leading to 401 token expiry error?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep, that's my understanding at least

useAuthStore.getState().setRefreshToken(refreshToken);
this.refreshToken = refreshToken;
this.axiosInstance.defaults.headers.common["Authorization"] = `Bearer ${token}`;
}
Expand All @@ -89,15 +93,27 @@ class ApiClient {
}

async refresh() {
const response = await this.axiosInstance.post<JsonValue>("/auth/refresh", {
refreshToken: this.refreshToken,
});
const resp = fromJson(RefreshTokenResponseSchema, response.data);
this.setTokens(resp.token, resp.refreshToken);
this.onTokenRefreshedEventEmitter.emit("tokenRefreshed", {
token: resp.token,
refreshToken: resp.refreshToken,
});
if (this.refreshPromise) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nice dedup logic. Yes this should address the multiple async calls to refresh, leading to multiple redundant token generation, and possible 401 bug across sessions.

This is actually rather tricky to test though since i think the bug is probabilistic? i believe for the current session, neither caller (your e.g. v2/chats/models and v2/chats/conversations) will bug out since auth-store uses last write valid token. It is only when caller_A and caller_B each calls refresh() and backend mongo receives a different set of token from auth-store.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yea, it's definitely not consistent. I was able to reproduce it a few (rare) times by lowering the expiration time of the token, and then switching between the settings and chat tabs. Although I am not 100% sure if it is the problem stated in #110

From my testing, my understanding is that it is possible for either caller_A or caller_B to bug out in the current session.

For example, caller_A and caller_B calls refresh() with token_A, there were a few possibilities that I encountered:

  1. caller_A refreshes to token_B and the API call is retried, caller_B fails to refresh with token_A because the backend is updated to token_B, and this API call is not retried and fails.
  2. caller_A receives a new token_B and caller_B receives a new token_C at the same time. Assuming the frontend receives token_C latest, it stores it in its ApiClient instance. However, in some rare occurrences, the backend stores token_B instead. This causes subsequent API calls to fail. Coincidentally, this happened again literally as I was typing this.
  3. Same as the above but the backend stores the same token as the frontend, so subsequent API calls succeed as per normal.

return this.refreshPromise;
}

this.refreshPromise = (async () => {
try {
const response = await this.axiosInstance.post<JsonValue>("/auth/refresh", {
refreshToken: this.refreshToken,
});
const resp = fromJson(RefreshTokenResponseSchema, response.data);
this.setTokens(resp.token, resp.refreshToken);
this.onTokenRefreshedEventEmitter.emit("tokenRefreshed", {
token: resp.token,
refreshToken: resp.refreshToken,
});
} finally {
this.refreshPromise = null;
}
})();

return this.refreshPromise;
}

private async requestWithRefresh(config: AxiosRequestConfig): Promise<JsonValue> {
Expand Down