From 488b0b71d6fa77f0e6fd4e449ab2c615b349918b Mon Sep 17 00:00:00 2001 From: Logan Ramos Date: Tue, 11 Nov 2025 10:29:56 -0500 Subject: [PATCH] Cleanup auto mode service (#1921) * Cleanup auto mode service * Update src/platform/endpoint/node/automodeService.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update src/platform/endpoint/node/automodeService.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../platform/endpoint/node/automodeService.ts | 273 +++++++----------- 1 file changed, 108 insertions(+), 165 deletions(-) diff --git a/extensions/copilot/src/platform/endpoint/node/automodeService.ts b/extensions/copilot/src/platform/endpoint/node/automodeService.ts index fd65a7afd4a..8ea765ce911 100644 --- a/extensions/copilot/src/platform/endpoint/node/automodeService.ts +++ b/extensions/copilot/src/platform/endpoint/node/automodeService.ts @@ -6,7 +6,7 @@ import { RequestType } from '@vscode/copilot-api'; import type { ChatRequest } from 'vscode'; import { createServiceIdentifier } from '../../../util/common/services'; -import { TaskSingler } from '../../../util/common/taskSingler'; +import { TimeoutTimer } from '../../../util/vs/base/common/async'; import { Disposable } from '../../../util/vs/base/common/lifecycle'; import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation'; import { IAuthenticationService } from '../../authentication/common/authentication'; @@ -24,141 +24,39 @@ interface AutoModeAPIResponse { session_token: string; } -/** - * Represents a cached auto mode token and the endpoint it maps to. - */ -interface CachedAutoToken { - readonly endpoint: IChatEndpoint; - readonly expiration: number; - readonly sessionToken: string; -} - -/** - * Holds the active and standby tokens for a conversation. - */ -interface ConversationCacheEntry { - active?: CachedAutoToken; - standby?: CachedAutoToken; -} - -export const IAutomodeService = createServiceIdentifier('IAutomodeService'); - -export interface IAutomodeService { - readonly _serviceBrand: undefined; - - resolveAutoModeEndpoint(chatRequest: ChatRequest | undefined, knownEndpoints: IChatEndpoint[]): Promise; -} - -export class AutomodeService extends Disposable implements IAutomodeService { - readonly _serviceBrand: undefined; - private readonly _autoModelCache: Map = new Map(); - private _reserveToken: CachedAutoToken | undefined; - private readonly _taskSingler = new TaskSingler(); - +class AutoModeTokenBank extends Disposable { + private _token: AutoModeAPIResponse | undefined; + private _fetchTokenPromise: Promise | undefined; + private _refreshTimer: TimeoutTimer; constructor( - @ICAPIClientService private readonly _capiClientService: ICAPIClientService, - @IAuthenticationService private readonly _authService: IAuthenticationService, - @ILogService private readonly _logService: ILogService, - @IInstantiationService private readonly _instantiationService: IInstantiationService, - @IExperimentationService private readonly _expService: IExperimentationService + public debugName: string, + private readonly _capiClientService: ICAPIClientService, + private readonly _authService: IAuthenticationService, + private readonly _logService: ILogService, + private readonly _expService: IExperimentationService ) { super(); - this._register(this._authService.onDidAuthenticationChange(() => { - this._autoModelCache.clear(); - this._reserveToken = undefined; - })); - this._serviceBrand = undefined; + this._refreshTimer = this._register(new TimeoutTimer()); + this._fetchTokenPromise = this._fetchToken(); } - /** - * Resolve an auto mode endpoint using a double-buffer strategy and a global reserve token. - */ - async resolveAutoModeEndpoint(chatRequest: ChatRequest | undefined, knownEndpoints: IChatEndpoint[]): Promise { - if (!knownEndpoints.length) { - throw new Error('No auto mode endpoints provided.'); + async getToken(): Promise { + if (!this._token) { + if (this._fetchTokenPromise) { + await this._fetchTokenPromise; + } else { + this._fetchTokenPromise = this._fetchToken(); + await this._fetchTokenPromise; + } } - - const conversationId = getConversationId(chatRequest); - const entry = this._autoModelCache.get(conversationId) ?? {}; - if (!this._autoModelCache.has(conversationId)) { - this._autoModelCache.set(conversationId, entry); + if (!this._token) { + throw new Error(`[${this.debugName}] Failed to fetch AutoMode token: token is undefined after fetch attempt.`); } - - this._pruneExpiredTokens(entry); - if (!entry.active && entry.standby) { - entry.active = entry.standby; - entry.standby = undefined; - } - - if (!entry.active) { - entry.active = await this._acquireActiveToken(conversationId, entry, knownEndpoints); - } - - if (!entry.standby || !this._isTokenValid(entry.standby) || this._isExpiringSoon(entry.standby) || this._isExpiringSoon(entry.active)) { - this._refreshStandbyInBackground(conversationId, entry, knownEndpoints); - } - - this._ensureReserveRefill(knownEndpoints); - return entry.active.endpoint; + return this._token; } - /** - * Acquire or refresh the reserve token so that a future conversation can respond instantly. - */ - private _ensureReserveRefill(knownEndpoints: IChatEndpoint[]): void { - if (this._isTokenValid(this._reserveToken)) { - return; - } - - void this._taskSingler.getOrCreate('reserve', () => this._fetchToken('reserve', undefined, knownEndpoints)) - .then(token => { - this._reserveToken = token; - }) - .catch(err => { - this._logService.error(`Failed to refresh reserve auto mode token: ${err instanceof Error ? err.message : String(err)}`); - }); - } - - /** - * Acquire the active token for a conversation, promoting the reserve if available. - */ - private async _acquireActiveToken(conversationId: string, entry: ConversationCacheEntry, knownEndpoints: IChatEndpoint[]): Promise { - if (this._isTokenValid(this._reserveToken)) { - const token = this._reserveToken; - this._reserveToken = undefined; - return token; - } - - const sessionHint = entry.standby?.sessionToken ?? entry.active?.sessionToken; - return this._taskSingler.getOrCreate(`active:${conversationId}`, () => this._fetchToken('active', sessionHint, knownEndpoints)); - } - - /** - * Start a background refresh to populate or update the standby token. - */ - private _refreshStandbyInBackground(conversationId: string, entrySnapshot: ConversationCacheEntry, knownEndpoints: IChatEndpoint[]): void { - const sessionHint = entrySnapshot.standby?.sessionToken ?? entrySnapshot.active?.sessionToken; - void this._taskSingler.getOrCreate(`standby:${conversationId}`, () => this._fetchToken('standby', sessionHint, knownEndpoints)) - .then(token => { - const entry = this._autoModelCache.get(conversationId); - if (!entry) { - return; - } - if (entry.active && entry.active.sessionToken === token.sessionToken) { - return; - } - entry.standby = token; - }) - .catch(err => { - this._logService.error(`Failed to refresh standby auto mode token for ${conversationId}: ${err instanceof Error ? err.message : String(err)}`); - }); - } - - /** - * Fetch a new token from the auto mode service. - */ - private async _fetchToken(debugName: string, sessionToken: string | undefined, knownEndpoints: IChatEndpoint[]): Promise { + private async _fetchToken(): Promise { const startTime = Date.now(); const authToken = (await this._authService.getCopilotToken()).token; @@ -166,8 +64,8 @@ export class AutomodeService extends Disposable implements IAutomodeService { 'Content-Type': 'application/json', 'Authorization': `Bearer ${authToken}` }; - if (sessionToken) { - headers['Copilot-Session-Token'] = sessionToken; + if (this._token) { + headers['Copilot-Session-Token'] = this._token.session_token; } const autoModeHint = this._expService.getTreatmentVariable('copilotchat.autoModelHint') || 'auto'; @@ -180,14 +78,88 @@ export class AutomodeService extends Disposable implements IAutomodeService { method: 'POST' }, { type: RequestType.AutoModels }); const data: AutoModeAPIResponse = await response.json() as AutoModeAPIResponse; - const selectedModel = knownEndpoints.find(e => e.model === data.selected_model) || knownEndpoints[0]; - const autoEndpoint = this._instantiationService.createInstance(AutoChatEndpoint, selectedModel, data.session_token, data.discounted_costs?.[selectedModel.model] || 0, this._calculateDiscountRange(data.discounted_costs)); - this._logService.trace(`Fetched auto model for ${debugName} in ${Date.now() - startTime}ms.`); - return { - endpoint: autoEndpoint, - expiration: data.expires_at * 1000, - sessionToken: data.session_token - }; + this._logService.trace(`Fetched auto model for ${this.debugName} in ${Date.now() - startTime}ms.`); + this._token = data; + // Trigger a refresh 5 minutes before expiration + this._refreshTimer.cancelAndSet(this._fetchToken.bind(this), (data.expires_at * 1000) - Date.now() - 5 * 60 * 1000); + this._fetchTokenPromise = undefined; + } + +} + +export const IAutomodeService = createServiceIdentifier('IAutomodeService'); + +export interface IAutomodeService { + readonly _serviceBrand: undefined; + + resolveAutoModeEndpoint(chatRequest: ChatRequest | undefined, knownEndpoints: IChatEndpoint[]): Promise; +} + +export class AutomodeService extends Disposable implements IAutomodeService { + readonly _serviceBrand: undefined; + private readonly _autoModelCache: Map = new Map(); + private _reserveToken: AutoModeTokenBank | undefined; + + constructor( + @ICAPIClientService private readonly _capiClientService: ICAPIClientService, + @IAuthenticationService private readonly _authService: IAuthenticationService, + @ILogService private readonly _logService: ILogService, + @IInstantiationService private readonly _instantiationService: IInstantiationService, + @IExperimentationService private readonly _expService: IExperimentationService + ) { + super(); + this._register(this._authService.onDidAuthenticationChange(() => { + for (const entry of this._autoModelCache.values()) { + entry.tokenBank.dispose(); + } + this._autoModelCache.clear(); + this._reserveToken?.dispose(); + this._reserveToken = new AutoModeTokenBank('reserve', this._capiClientService, this._authService, this._logService, this._expService); + })); + this._serviceBrand = undefined; + } + + override dispose(): void { + for (const entry of this._autoModelCache.values()) { + entry.tokenBank.dispose(); + } + this._autoModelCache.clear(); + this._reserveToken?.dispose(); + super.dispose(); + } + + /** + * Resolve an auto mode endpoint using a double-buffer strategy and a global reserve token. + */ + async resolveAutoModeEndpoint(chatRequest: ChatRequest | undefined, knownEndpoints: IChatEndpoint[]): Promise { + if (!knownEndpoints.length) { + throw new Error('No auto mode endpoints provided.'); + } + + const conversationId = getConversationId(chatRequest); + const entry = this._autoModelCache.get(conversationId); + if (entry) { + const entryToken = await entry.tokenBank.getToken(); + if (entry.endpoint.model !== entryToken.selected_model) { + // Model changed during a token refresh -> map to new endpoint + const newModel = knownEndpoints.find(e => e.model === entryToken.selected_model) || knownEndpoints[0]; + entry.endpoint = this._instantiationService.createInstance(AutoChatEndpoint, newModel, entryToken.session_token, entryToken.discounted_costs?.[newModel.model] || 0, this._calculateDiscountRange(entryToken.discounted_costs)); + } + return entry.endpoint; + } + + // No entry yet -> Promote reserve token to active and repopulate reserve + const reserveTokenBank = this._reserveToken || new AutoModeTokenBank('reserve', this._capiClientService, this._authService, this._logService, this._expService); + this._reserveToken = new AutoModeTokenBank('reserve', this._capiClientService, this._authService, this._logService, this._expService); + + // Update the debug name so logs are properly associating this token with the right conversation id now + reserveTokenBank.debugName = conversationId; + + const reserveToken = await reserveTokenBank.getToken(); + const selectedModel = knownEndpoints.find(e => e.model === reserveToken.selected_model) || knownEndpoints[0]; + const autoEndpoint = this._instantiationService.createInstance(AutoChatEndpoint, selectedModel, reserveToken.session_token, reserveToken.discounted_costs?.[selectedModel.model] || 0, this._calculateDiscountRange(reserveToken.discounted_costs)); + this._autoModelCache.set(conversationId, { endpoint: autoEndpoint, tokenBank: reserveTokenBank }); + return autoEndpoint; } private _calculateDiscountRange(discounts: Record | undefined): { low: number; high: number } { @@ -209,35 +181,6 @@ export class AutomodeService extends Disposable implements IAutomodeService { } return hasValues ? { low, high } : { low: 0, high: 0 }; } - - /** - * Remove expired tokens so they are not considered during promotion. - */ - private _pruneExpiredTokens(entry: ConversationCacheEntry): void { - if (entry.active && !this._isTokenValid(entry.active)) { - entry.active = undefined; - } - if (entry.standby && !this._isTokenValid(entry.standby)) { - entry.standby = undefined; - } - } - - /** - * Determine whether a token is still valid. - */ - private _isTokenValid(token: CachedAutoToken | undefined): token is CachedAutoToken { - return !!token && token.expiration > Date.now(); - } - - /** - * Determine whether a token should be refreshed soon. - */ - private _isExpiringSoon(token: CachedAutoToken | undefined): boolean { - if (!token) { - return false; - } - return token.expiration - Date.now() <= 5 * 60 * 1000; - } } /**