Skip to content

Commit 88d8741

Browse files
sbattensbatten
authored
Fix client-side Copilot token handling and add tests (#331737)
* Agent Host changes for sbatten/agents/client-side-copilot-token-fixes-a26d0eed * Refactor fetchedValue utility and improve tests for better coverage
1 parent 680caf2 commit 88d8741

11 files changed

Lines changed: 461 additions & 40 deletions

File tree

‎extensions/copilot/src/extension/completions-core/vscode-node/lib/src/auth/copilotTokenManager.ts‎

Lines changed: 23 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
import { IAuthenticationService } from '../../../../../../platform/authentication/common/authentication';
77
import { CopilotToken } from '../../../../../../platform/authentication/common/copilotToken';
8+
import { FetchedValue } from '../../../../../../shared-fetch-utils/common/fetchedValue';
89
import { createServiceIdentifier } from '../../../../../../util/common/services';
910
import { ThrottledDelayer } from '../../../../../../util/vs/base/common/async';
1011
import { Disposable } from '../../../../../../util/vs/base/common/lifecycle';
@@ -22,11 +23,14 @@ export interface ICompletionsCopilotTokenManager {
2223

2324
export class CopilotTokenManagerImpl extends Disposable implements ICompletionsCopilotTokenManager {
2425
declare _serviceBrand: undefined;
25-
private tokenRefetcher = new ThrottledDelayer(5_000);
26-
private _token: CopilotToken | undefined;
26+
private readonly tokenRefetcher: ThrottledDelayer<CopilotToken>;
27+
private readonly tokenValue: FetchedValue<CopilotToken>;
28+
2729
get token() {
28-
void this.tokenRefetcher.trigger(() => this.updateCachedToken());
29-
return this._token;
30+
void this.tokenRefetcher.trigger(() => this.tokenValue.resolve()).catch(() => {
31+
// Foreground getToken calls surface the cached error.
32+
});
33+
return this.tokenValue.value;
3034
}
3135

3236
constructor(
@@ -35,8 +39,15 @@ export class CopilotTokenManagerImpl extends Disposable implements ICompletionsC
3539
) {
3640
super();
3741

38-
this.updateCachedToken();
39-
this._register(this.authenticationService.onDidCopilotTokenChange(() => this.updateCachedToken()));
42+
this.tokenRefetcher = this._register(new ThrottledDelayer(5_000));
43+
this.tokenValue = this._register(new FetchedValue({
44+
fetch: () => this.authenticationService.getCopilotToken(),
45+
isStale: () => true,
46+
getRetryAfterMs: () => 5_000,
47+
}));
48+
49+
this.resolveInBackground();
50+
this._register(this.authenticationService.onDidCopilotTokenChange(() => this.resolveInBackground(true)));
4051
}
4152

4253
/**
@@ -54,15 +65,17 @@ export class CopilotTokenManagerImpl extends Disposable implements ICompletionsC
5465
}
5566

5667
async getToken(): Promise<CopilotToken> {
57-
return this.updateCachedToken();
68+
return this.tokenValue.resolve();
5869
}
5970

60-
private async updateCachedToken(): Promise<CopilotToken> {
61-
this._token = await this.authenticationService.getCopilotToken();
62-
return this._token;
71+
private resolveInBackground(force?: boolean): void {
72+
void this.tokenValue.resolve(force).catch(() => {
73+
// Foreground getToken calls surface the cached error.
74+
});
6375
}
6476

6577
resetToken(httpError?: number): void {
78+
this.tokenValue.invalidate();
6679
this.authenticationService.resetCopilotToken();
6780
}
6881

‎extensions/copilot/src/extension/completions-core/vscode-node/lib/src/auth/test/copilotTokenManager.spec.ts‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
7+
import { IAuthenticationService } from '../../../../../../../platform/authentication/common/authentication';
8+
import { StaticGitHubAuthenticationService } from '../../../../../../../platform/authentication/common/staticGitHubAuthenticationService';
9+
import { CopilotToken, createTestExtendedTokenInfo } from '../../../../../../../platform/authentication/common/copilotToken';
10+
import { ICopilotTokenManager } from '../../../../../../../platform/authentication/common/copilotTokenManager';
11+
import { ICopilotTokenStore } from '../../../../../../../platform/authentication/common/copilotTokenStore';
12+
import { IConfigurationService } from '../../../../../../../platform/configuration/common/configurationService';
13+
import { ILogService } from '../../../../../../../platform/log/common/logService';
14+
import { createPlatformServices, ITestingServicesAccessor } from '../../../../../../../platform/test/node/services';
15+
import { FetchBlockedError } from '../../../../../../../shared-fetch-utils/common/fetchTypes';
16+
import { Event } from '../../../../../../../util/vs/base/common/event';
17+
import { DisposableStore } from '../../../../../../../util/vs/base/common/lifecycle';
18+
import { CopilotTokenManagerImpl } from '../copilotTokenManager';
19+
20+
describe('CopilotTokenManagerImpl', () => {
21+
let accessor: ITestingServicesAccessor;
22+
let disposables: DisposableStore;
23+
24+
beforeEach(() => {
25+
vi.useFakeTimers();
26+
vi.setSystemTime(100);
27+
disposables = new DisposableStore();
28+
accessor = disposables.add(createPlatformServices().createTestingAccessor());
29+
});
30+
31+
afterEach(() => {
32+
disposables.dispose();
33+
vi.useRealTimers();
34+
});
35+
36+
it('caches ordinary failures for five seconds', async () => {
37+
const tokenManager = new FailingCopilotTokenManager(() => new Error('network failure'));
38+
const manager = createManager(tokenManager);
39+
40+
await expect(manager.getToken()).rejects.toThrow('network failure');
41+
await expect(manager.getToken()).rejects.toThrow('network failure');
42+
expect(tokenManager.calls).toBe(1);
43+
44+
vi.advanceTimersByTime(4_999);
45+
await expect(manager.getToken()).rejects.toThrow('network failure');
46+
expect(tokenManager.calls).toBe(1);
47+
48+
vi.advanceTimersByTime(1);
49+
await expect(manager.getToken()).rejects.toThrow('network failure');
50+
expect(tokenManager.calls).toBe(2);
51+
});
52+
53+
it('prefers a server retry delay over the fallback cooldown', async () => {
54+
const tokenManager = new FailingCopilotTokenManager(() => new FetchBlockedError('rate limited', 30_000));
55+
const manager = createManager(tokenManager);
56+
57+
await expect(manager.getToken()).rejects.toThrow('rate limited');
58+
vi.advanceTimersByTime(5_000);
59+
await expect(manager.getToken()).rejects.toThrow('rate limited');
60+
expect(tokenManager.calls).toBe(1);
61+
62+
vi.advanceTimersByTime(25_000);
63+
await expect(manager.getToken()).rejects.toThrow('rate limited');
64+
expect(tokenManager.calls).toBe(2);
65+
});
66+
67+
it('foreground calls still fail after a successful token fetch', async () => {
68+
const token = new CopilotToken(createTestExtendedTokenInfo({ token: 'tid=success' }));
69+
const tokenManager = new ScriptedCopilotTokenManager([
70+
token,
71+
new Error('signed out'),
72+
]);
73+
const manager = createManager(tokenManager);
74+
75+
await expect(manager.getToken()).resolves.toBe(token);
76+
await expect(manager.getToken()).rejects.toThrow('signed out');
77+
await expect(manager.primeToken()).resolves.toBe(false);
78+
79+
expect(manager.token).toBe(token);
80+
expect(tokenManager.calls).toBe(2);
81+
});
82+
83+
function createManager(tokenManager: ICopilotTokenManager): CopilotTokenManagerImpl {
84+
const authenticationService: IAuthenticationService = disposables.add(new StaticGitHubAuthenticationService(
85+
() => 'github-token',
86+
accessor.get(ILogService),
87+
accessor.get(ICopilotTokenStore),
88+
tokenManager,
89+
accessor.get(IConfigurationService),
90+
));
91+
return disposables.add(new CopilotTokenManagerImpl(false, authenticationService));
92+
}
93+
});
94+
95+
class FailingCopilotTokenManager implements ICopilotTokenManager {
96+
declare readonly _serviceBrand: undefined;
97+
readonly onDidCopilotTokenRefresh = Event.None;
98+
calls = 0;
99+
100+
constructor(private readonly createError: () => Error) { }
101+
102+
async getCopilotToken(): Promise<CopilotToken> {
103+
this.calls++;
104+
throw this.createError();
105+
}
106+
107+
resetCopilotToken(): void { }
108+
}
109+
110+
class ScriptedCopilotTokenManager implements ICopilotTokenManager {
111+
declare readonly _serviceBrand: undefined;
112+
readonly onDidCopilotTokenRefresh = Event.None;
113+
calls = 0;
114+
115+
constructor(private readonly results: Array<CopilotToken | Error>) { }
116+
117+
async getCopilotToken(): Promise<CopilotToken> {
118+
this.calls++;
119+
const result = this.results.shift();
120+
if (!result) {
121+
throw new Error('No scripted token result');
122+
}
123+
if (result instanceof Error) {
124+
throw result;
125+
}
126+
return result;
127+
}
128+
129+
resetCopilotToken(): void { }
130+
}

‎extensions/copilot/src/platform/authentication/common/authentication.ts‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,6 @@ export abstract class BaseAuthenticationService extends Disposable implements IA
278278

279279
//#region Copilot Token
280280

281-
private _copilotTokenError: Error | undefined;
282281
get copilotToken(): CopilotToken | undefined {
283282
return this._tokenStore.copilotToken;
284283
}
@@ -287,24 +286,16 @@ export abstract class BaseAuthenticationService extends Disposable implements IA
287286
const tokenBefore = this._tokenStore.copilotToken;
288287
const token = await this._tokenManager.getCopilotToken(force);
289288
this._tokenStore.copilotToken = token;
290-
this._copilotTokenError = undefined;
291289
if (tokenBefore?.token !== token.token) {
292290
this.fireCopilotTokenChange('getCopilotToken');
293291
}
294292
return token;
295293
} catch (afterError) {
296294
const tokenBefore = this._tokenStore.copilotToken;
297295
this._tokenStore.copilotToken = undefined;
298-
const beforeError = this._copilotTokenError;
299-
this._copilotTokenError = afterError;
300296
if (tokenBefore) {
301297
// Had a valid token before, now errored — token value changed to undefined
302298
this.fireCopilotTokenChange('getCopilotToken token lost');
303-
} else if (beforeError && afterError && beforeError.message !== afterError.message) {
304-
// Still can't get a Copilot Token, but the error has changed.
305-
// I.e. They go from being not signed in (no copilot token can be minted)
306-
// to an account that doesn't have a valid subscription (no copilot token can be minted).
307-
this.fireCopilotTokenChange('getCopilotToken error change');
308299
}
309300
throw afterError;
310301
}

‎extensions/copilot/src/platform/authentication/common/copilotToken.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -633,6 +633,8 @@ export type SuccessNotificationId =
633633

634634
export type TokenError = {
635635
reason: TokenErrorReason;
636+
/** Milliseconds the caller should wait before retrying a rate-limited request. */
637+
retryAfterMs?: number;
636638
notification_id?: TokenErrorNotificationId | string;
637639
message?: string;
638640
/** URL for action button to help user resolve the error. */

‎extensions/copilot/src/platform/authentication/node/copilotTokenManager.ts‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
*--------------------------------------------------------------------------------------------*/
55

66
import { RequestType } from '@vscode/copilot-api';
7+
import { retryAfterFromRateLimitHeaders } from '../../../shared-fetch-utils/common/middleware/rateLimitBackoffMiddleware';
78
import { Emitter } from '../../../util/vs/base/common/event';
89
import { Disposable, toDisposable } from '../../../util/vs/base/common/lifecycle';
910
import { SyncDescriptor } from '../../../util/vs/platform/instantiation/common/descriptors';
@@ -28,6 +29,7 @@ type FetchTokenResult = {
2829
ok: boolean;
2930
status: number;
3031
statusText: string;
32+
retryAfterMs?: number;
3133
} & (
3234
// success
3335
| { body: TokenEnvelope; kind: 'token' }
@@ -180,6 +182,10 @@ export abstract class BaseCopilotTokenManager extends Disposable implements ICop
180182
// Handle HTTP errors
181183
if (!result.ok) {
182184
this._logService.warn(`Failed to get copilot token due to status ${result.status} ${result.statusText}`);
185+
if (result.status === 429) {
186+
this._telemetryService.sendGHTelemetryErrorEvent('auth.rate_limited');
187+
return { kind: 'failure', reason: 'RateLimited', retryAfterMs: result.retryAfterMs };
188+
}
183189
const data = TelemetryData.createAndMarkAsIssued({
184190
status: result.status.toString(),
185191
status_text: result.statusText,
@@ -206,7 +212,7 @@ export abstract class BaseCopilotTokenManager extends Disposable implements ICop
206212
if (result.body.message?.startsWith('API rate limit exceeded')) {
207213
this._logService.warn('Failed to get copilot token due to exceeding API rate limit');
208214
this._telemetryService.sendGHTelemetryErrorEvent('auth.rate_limited');
209-
return { kind: 'failure', reason: 'RateLimited' };
215+
return { kind: 'failure', reason: 'RateLimited', retryAfterMs: result.retryAfterMs };
210216
}
211217
this._logService.warn(`Failed to get copilot token due to: ${result.body.message}`);
212218
return { kind: 'failure', reason: 'NotAuthorized' };
@@ -290,7 +296,12 @@ export abstract class BaseCopilotTokenManager extends Disposable implements ICop
290296
* Returns a structured result with HTTP status and validated body.
291297
*/
292298
private async parseTokenResponse(response: Response): Promise<FetchTokenResult> {
293-
const httpInfo = { ok: response.ok, status: response.status, statusText: response.statusText };
299+
const httpInfo = {
300+
ok: response.ok,
301+
status: response.status,
302+
statusText: response.statusText,
303+
retryAfterMs: retryAfterFromRateLimitHeaders(response.headers),
304+
};
294305

295306
let parsed: unknown;
296307
try {

‎extensions/copilot/src/platform/authentication/test/node/authentication.spec.ts‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import { ITelemetryService } from '../../../telemetry/common/telemetry';
1616
import { createPlatformServices } from '../../../test/node/services';
1717
import { StaticGitHubAuthenticationService } from '../../common/staticGitHubAuthenticationService';
1818
import { CopilotToken, createTestExtendedTokenInfo } from '../../common/copilotToken';
19+
import { ICopilotTokenManager } from '../../common/copilotTokenManager';
1920
import { ICopilotTokenStore } from '../../common/copilotTokenStore';
2021
import { FixedCopilotTokenManager } from '../../node/copilotTokenManager';
2122

@@ -109,4 +110,66 @@ suite('AuthenticationService', function () {
109110
await promise;
110111
expect(authenticationService.copilotToken?.token).toBe(newToken);
111112
});
113+
114+
test('Does not emit onDidCopilotTokenChange when token errors change', async () => {
115+
const accessor = disposables.add(createPlatformServices().createTestingAccessor());
116+
const failingTokenManager = new ScriptedCopilotTokenManager([
117+
new Error('first failure'),
118+
new Error('second failure'),
119+
]);
120+
const service = disposables.add(new StaticGitHubAuthenticationService(
121+
() => testToken,
122+
accessor.get(ILogService),
123+
accessor.get(ICopilotTokenStore),
124+
failingTokenManager,
125+
accessor.get(IConfigurationService),
126+
));
127+
const tokenChangeSpy = vi.fn();
128+
service.onDidCopilotTokenChange(tokenChangeSpy);
129+
130+
await expect(service.getCopilotToken()).rejects.toThrow('first failure');
131+
await expect(service.getCopilotToken()).rejects.toThrow('second failure');
132+
133+
expect(tokenChangeSpy).not.toHaveBeenCalled();
134+
});
135+
136+
test('Emits onDidCopilotTokenChange when a token is gained and lost', async () => {
137+
const accessor = disposables.add(createPlatformServices().createTestingAccessor());
138+
const token = new CopilotToken(createTestExtendedTokenInfo({ token: 'tid=scripted' }));
139+
const tokenManager = new ScriptedCopilotTokenManager([token, new Error('token lost')]);
140+
const service = disposables.add(new StaticGitHubAuthenticationService(
141+
() => testToken,
142+
accessor.get(ILogService),
143+
accessor.get(ICopilotTokenStore),
144+
tokenManager,
145+
accessor.get(IConfigurationService),
146+
));
147+
const observedTokens: Array<string | undefined> = [];
148+
service.onDidCopilotTokenChange(() => observedTokens.push(service.copilotToken?.token));
149+
150+
await service.getCopilotToken();
151+
await expect(service.getCopilotToken()).rejects.toThrow('token lost');
152+
153+
expect(observedTokens).toEqual(['tid=scripted', undefined]);
154+
});
112155
});
156+
157+
class ScriptedCopilotTokenManager implements ICopilotTokenManager {
158+
declare readonly _serviceBrand: undefined;
159+
readonly onDidCopilotTokenRefresh = Event.None;
160+
161+
constructor(private readonly results: Array<CopilotToken | Error>) { }
162+
163+
async getCopilotToken(): Promise<CopilotToken> {
164+
const result = this.results.shift();
165+
if (!result) {
166+
throw new Error('No scripted token result');
167+
}
168+
if (result instanceof Error) {
169+
throw result;
170+
}
171+
return result;
172+
}
173+
174+
resetCopilotToken(): void { }
175+
}

0 commit comments

Comments
 (0)