mirror of
https://github.com/microsoft/vscode.git
synced 2026-09-27 10:37:38 +01:00
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>
This commit is contained in:
Logan Ramos
committed by
GitHub
GitHub
co-authored by
Copilot
Copilot
parent
7945d35c06
commit
488b0b71d6
@@ -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>('IAutomodeService');
|
||||
|
||||
export interface IAutomodeService {
|
||||
readonly _serviceBrand: undefined;
|
||||
|
||||
resolveAutoModeEndpoint(chatRequest: ChatRequest | undefined, knownEndpoints: IChatEndpoint[]): Promise<IChatEndpoint>;
|
||||
}
|
||||
|
||||
export class AutomodeService extends Disposable implements IAutomodeService {
|
||||
readonly _serviceBrand: undefined;
|
||||
private readonly _autoModelCache: Map<string, ConversationCacheEntry> = new Map();
|
||||
private _reserveToken: CachedAutoToken | undefined;
|
||||
private readonly _taskSingler = new TaskSingler<CachedAutoToken>();
|
||||
|
||||
class AutoModeTokenBank extends Disposable {
|
||||
private _token: AutoModeAPIResponse | undefined;
|
||||
private _fetchTokenPromise: Promise<void> | 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<IChatEndpoint> {
|
||||
if (!knownEndpoints.length) {
|
||||
throw new Error('No auto mode endpoints provided.');
|
||||
async getToken(): Promise<AutoModeAPIResponse> {
|
||||
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<CachedAutoToken> {
|
||||
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<CachedAutoToken> {
|
||||
private async _fetchToken(): Promise<void> {
|
||||
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<string>('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>('IAutomodeService');
|
||||
|
||||
export interface IAutomodeService {
|
||||
readonly _serviceBrand: undefined;
|
||||
|
||||
resolveAutoModeEndpoint(chatRequest: ChatRequest | undefined, knownEndpoints: IChatEndpoint[]): Promise<IChatEndpoint>;
|
||||
}
|
||||
|
||||
export class AutomodeService extends Disposable implements IAutomodeService {
|
||||
readonly _serviceBrand: undefined;
|
||||
private readonly _autoModelCache: Map<string, { endpoint: IChatEndpoint; tokenBank: AutoModeTokenBank }> = 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<IChatEndpoint> {
|
||||
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<string, number> | 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;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
Reference in New Issue
Block a user