Skip to content

Commit 76f089f

Browse files
committed
Remove refreshIfAlmostExpired and replace it with smarter timers
1 parent 3570c64 commit 76f089f

File tree

1 file changed

+89
-34
lines changed

1 file changed

+89
-34
lines changed

src/oauth/sessionManager.ts

Lines changed: 89 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ export class OAuthSessionManager implements vscode.Disposable {
8181
private refreshPromise: Promise<TokenResponse> | null = null;
8282
private lastRefreshAttempt = 0;
8383
private refreshTimer: NodeJS.Timeout | undefined;
84+
private tokenChangeListener: vscode.Disposable | undefined;
8485

8586
private pendingAuthReject: ((reason: Error) => void) | undefined;
8687

@@ -99,7 +100,8 @@ export class OAuthSessionManager implements vscode.Disposable {
99100
container.getLoginCoordinator(),
100101
extensionId,
101102
);
102-
manager.scheduleBackgroundRefresh();
103+
manager.setupTokenListener();
104+
manager.scheduleNextRefresh();
103105
return manager;
104106
}
105107

@@ -162,30 +164,96 @@ export class OAuthSessionManager implements vscode.Disposable {
162164
}
163165

164166
/**
165-
* Clear refresh-related state.
167+
* Clear all refresh-related state: in-flight promise, throttle, timer, and listener.
166168
*/
167169
private clearRefreshState(): void {
168170
this.refreshPromise = null;
169171
this.lastRefreshAttempt = 0;
172+
if (this.refreshTimer) {
173+
clearTimeout(this.refreshTimer);
174+
this.refreshTimer = undefined;
175+
}
176+
this.tokenChangeListener?.dispose();
177+
this.tokenChangeListener = undefined;
178+
}
179+
180+
/**
181+
* Setup listener for token changes. Disposes existing listener first.
182+
* Reschedules refresh when tokens change (e.g., from another window).
183+
*/
184+
private setupTokenListener(): void {
185+
this.tokenChangeListener?.dispose();
186+
this.tokenChangeListener = undefined;
187+
188+
if (!this.deployment) {
189+
return;
190+
}
191+
192+
this.tokenChangeListener = this.secretsManager.onDidChangeOAuthTokens(
193+
this.deployment.safeHostname,
194+
() => this.scheduleNextRefresh(),
195+
);
170196
}
171197

172198
/**
173-
* Schedule the next background token refresh check.
174-
* Only schedules the next check after the current one completes.
199+
* Schedule the next token refresh based on expiry time.
200+
* - Far from expiry: schedule once at threshold
201+
* - Near/past expiry: attempt refresh immediately
175202
*/
176-
private scheduleBackgroundRefresh(): void {
203+
private scheduleNextRefresh(): void {
177204
if (this.refreshTimer) {
178205
clearTimeout(this.refreshTimer);
206+
this.refreshTimer = undefined;
179207
}
180208

181-
this.refreshTimer = setTimeout(async () => {
182-
try {
183-
await this.refreshIfAlmostExpired();
184-
} catch (error) {
185-
this.logger.warn("Background token refresh failed:", error);
186-
}
187-
this.scheduleBackgroundRefresh();
188-
}, BACKGROUND_REFRESH_INTERVAL_MS);
209+
this.getStoredTokens()
210+
.then((storedTokens) => {
211+
if (!storedTokens?.refresh_token) {
212+
return;
213+
}
214+
215+
const now = Date.now();
216+
const timeUntilExpiry = storedTokens.expiry_timestamp - now;
217+
218+
if (timeUntilExpiry <= TOKEN_REFRESH_THRESHOLD_MS) {
219+
// Within threshold or expired, attempt refresh now
220+
this.attemptRefreshWithRetry();
221+
} else {
222+
// Schedule for when we reach the threshold
223+
const delay = timeUntilExpiry - TOKEN_REFRESH_THRESHOLD_MS;
224+
this.logger.debug(
225+
`Scheduling token refresh in ${Math.round(delay / 1000 / 60)} minutes`,
226+
);
227+
this.refreshTimer = setTimeout(
228+
() => this.attemptRefreshWithRetry(),
229+
delay,
230+
);
231+
}
232+
})
233+
.catch((error) => {
234+
this.logger.warn("Failed to schedule token refresh:", error);
235+
});
236+
}
237+
238+
/**
239+
* Attempt refresh, falling back to polling on failure.
240+
*/
241+
private attemptRefreshWithRetry(): void {
242+
this.refreshTimer = undefined;
243+
244+
this.refreshToken()
245+
.then(() => {
246+
// Success - scheduleNextRefresh will be triggered by token change listener
247+
this.logger.debug("Background token refresh succeeded");
248+
})
249+
.catch((error) => {
250+
this.logger.warn("Background token refresh failed, will retry:", error);
251+
// Fall back to polling until successful
252+
this.refreshTimer = setTimeout(
253+
() => this.attemptRefreshWithRetry(),
254+
BACKGROUND_REFRESH_INTERVAL_MS,
255+
);
256+
});
189257
}
190258

191259
/**
@@ -327,18 +395,19 @@ export class OAuthSessionManager implements vscode.Disposable {
327395
this.deployment = deployment;
328396
this.clearRefreshState();
329397

330-
// Refresh tokens if needed to prevent 401s
331-
if (await this.isTokenExpired()) {
398+
// Block on refresh if token is expired to ensure valid state for callers
399+
const storedTokens = await this.getStoredTokens();
400+
if (storedTokens && Date.now() >= storedTokens.expiry_timestamp) {
332401
try {
333402
await this.refreshToken();
334403
} catch (error) {
335404
this.logger.warn("Token refresh failed (expired):", error);
336405
}
337-
} else {
338-
this.refreshIfAlmostExpired().catch((error) => {
339-
this.logger.warn("Background token refresh failed:", error);
340-
});
341406
}
407+
408+
// Schedule after blocking refresh to avoid concurrent attempts
409+
this.setupTokenListener();
410+
this.scheduleNextRefresh();
342411
}
343412

344413
public clearDeployment(): void {
@@ -375,6 +444,7 @@ export class OAuthSessionManager implements vscode.Disposable {
375444
});
376445
this.clearRefreshState();
377446
this.deployment = deployment;
447+
this.setupTokenListener();
378448
}
379449

380450
const client = CoderApi.create(deployment.url, undefined, this.logger);
@@ -735,17 +805,6 @@ export class OAuthSessionManager implements vscode.Disposable {
735805
return timeUntilExpiry < TOKEN_REFRESH_THRESHOLD_MS;
736806
}
737807

738-
/**
739-
* Check if token is expired.
740-
*/
741-
private async isTokenExpired(): Promise<boolean> {
742-
const storedTokens = await this.getStoredTokens();
743-
if (!storedTokens) {
744-
return false;
745-
}
746-
return Date.now() >= storedTokens.expiry_timestamp;
747-
}
748-
749808
public async revokeRefreshToken(): Promise<void> {
750809
const storedTokens = await this.getStoredTokens();
751810
if (!storedTokens?.refresh_token) {
@@ -868,10 +927,6 @@ export class OAuthSessionManager implements vscode.Disposable {
868927
* Clears all in-memory state.
869928
*/
870929
public dispose(): void {
871-
if (this.refreshTimer) {
872-
clearTimeout(this.refreshTimer);
873-
this.refreshTimer = undefined;
874-
}
875930
if (this.pendingAuthReject) {
876931
this.pendingAuthReject(new Error("OAuth session manager disposed"));
877932
}

0 commit comments

Comments
 (0)