Add a new fetchedValue concept (#4943)

* Implement fetchedValue

* Implement a compostible fetcher

* Keep cache hot

* Create a capi client fetched value

* Revert capi client change

* Fix tests

* Cleanup the middlewares to be a bit smarter

* cleanup stream cloning

* Update parse

* Fix window active cloning

* Fix tests

* Fix docs

* Update docs

Author: lramos15

eslint.config.mjs

@@ -245,6 +245,10 @@ export default tseslint.config(
 							target: './src/',
 							from: './test/'
 						},
+						{
+							target: './src/shared-fetch-utils',
+							from: ['./src/extension', './src/platform', './src/util', './src/lib']
+						},
 						{
 							target: './src/util',
 							from: ['./src/platform', './src/extension']

src/platform/endpoint/node/automodeService.ts

@@ -5,15 +5,16 @@
 
 import { RequestType } from '@vscode/copilot-api';
 import type { ChatRequest } from 'vscode';
+import { FetchedValue } from '../../../shared-fetch-utils/common/fetchedValue';
 import { createServiceIdentifier } from '../../../util/common/services';
-import { TimeoutTimer } from '../../../util/vs/base/common/async';
 import { Disposable, DisposableMap } from '../../../util/vs/base/common/lifecycle';
 import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation';
 import { ChatLocation } from '../../../vscodeTypes';
 import { IAuthenticationService } from '../../authentication/common/authentication';
 import { ConfigKey, IConfigurationService } from '../../configuration/common/configurationService';
 import { IEnvService } from '../../env/common/envService';
 import { ILogService } from '../../log/common/logService';
+import { createCapiClientFetchedValue } from '../../networking/common/capiClientFetchedValue';
 import { isAbortError } from '../../networking/common/fetcherService';
 import { IChatEndpoint } from '../../networking/common/networking';
 import { IRequestLogger } from '../../requestLogger/node/requestLogger';
@@ -41,108 +42,60 @@ interface AutoModelCacheEntry {
 }
 
 class AutoModeTokenBank extends Disposable {
-	private _token: AutoModeAPIResponse | undefined;
-	private _fetchTokenPromise: Promise<void> | undefined;
-	private _refreshTimer: TimeoutTimer;
+	private readonly _fetchedValue: FetchedValue<AutoModeAPIResponse>;
 	private _usedSinceLastFetch = false;
 
 	constructor(
 		public debugName: string,
-		private readonly _location: ChatLocation,
-		private readonly _capiClientService: ICAPIClientService,
-		private readonly _authService: IAuthenticationService,
-		private readonly _logService: ILogService,
-		private readonly _expService: IExperimentationService,
-		private readonly _envService: IEnvService
+		location: ChatLocation,
+		capiClientService: ICAPIClientService,
+		authService: IAuthenticationService,
+		_logService: ILogService,
+		expService: IExperimentationService,
+		envService: IEnvService,
 	) {
 		super();
-		this._refreshTimer = this._register(new TimeoutTimer());
-		this._register(this._envService.onDidChangeWindowState((state) => {
-			if (state.active && this._usedSinceLastFetch && (!this._token || this._token.expires_at * 1000 - Date.now() < 5 * 60 * 1000)) {
-				// Window is active again, fetch a new token if it's expiring soon or we don't have one
-				this._fetchTokenPromise = this._fetchToken();
-			}
+
+		const expName = location === ChatLocation.Editor
+			? 'copilotchat.autoModelHint.editor'
+			: 'copilotchat.autoModelHint';
+
+		this._fetchedValue = this._register(createCapiClientFetchedValue<AutoModeAPIResponse>(capiClientService, envService, {
+			request: async () => {
+				const authToken = (await authService.getCopilotToken()).token;
+				const autoModeHint = expService.getTreatmentVariable<string>(expName) || 'auto';
+				return {
+					headers: {
+						'Content-Type': 'application/json',
+						'Authorization': `Bearer ${authToken}`,
+					},
+					method: 'POST' as const,
+					json: { auto_mode: { model_hints: [autoModeHint] } },
+				};
+			},
+			requestMetadata: { type: RequestType.AutoModels },
+			parseResponse: async (res) => {
+				if (res.status < 200 || res.status >= 300) {
+					const text = await res.text().catch(() => '');
+					throw new Error(`AutoMode token response status: ${res.status}${text ? `, body: ${text}` : ''}`);
+				}
+				const data = await res.json() as AutoModeAPIResponse;
+				this._usedSinceLastFetch = false;
+				return data;
+			},
+			isStale: (token) => {
+				if (!this._usedSinceLastFetch) {
+					return false;
+				}
+				return token.expires_at * 1000 - Date.now() < 5 * 60 * 1000;
+			},
+			keepCacheHot: true,
 		}));
-		this._fetchTokenPromise = this._fetchToken();
 	}
 
 	async getToken(): Promise<AutoModeAPIResponse> {
-		if (!this._token) {
-			if (this._fetchTokenPromise) {
-				await this._fetchTokenPromise;
-			}
-			// If we still don't have a token (e.g., the awaited promise returned nothing), force a new fetch
-			if (!this._token) {
-				this._fetchTokenPromise = this._fetchToken(true);
-				await this._fetchTokenPromise;
-			}
-		}
-		if (!this._token) {
-			throw new Error(`[${this.debugName}] Failed to fetch AutoMode token: token is undefined after fetch attempt.`);
-		}
 		this._usedSinceLastFetch = true;
-		return this._token;
-	}
-
-
-	private async _fetchToken(force?: boolean): Promise<void> {
-		// If the window isn't active we will skip fetching to save network calls
-		// We will fetch again when the window becomes active
-		if (!this._envService.isActive && !force) {
-			return;
-		}
-		const startTime = Date.now();
-
-		try {
-			const authToken = (await this._authService.getCopilotToken()).token;
-			const headers: Record<string, string> = {
-				'Content-Type': 'application/json',
-				'Authorization': `Bearer ${authToken}`
-			};
-
-			const expName = this._location === ChatLocation.Editor
-				? 'copilotchat.autoModelHint.editor'
-				: 'copilotchat.autoModelHint';
-
-			const autoModeHint = this._expService.getTreatmentVariable<string>(expName) || 'auto';
-
-			const response = await this._capiClientService.makeRequest<Response>({
-				json: {
-					'auto_mode': { 'model_hints': [autoModeHint] }
-				},
-				headers,
-				method: 'POST'
-			}, { type: RequestType.AutoModels });
-			if (!response.ok) {
-				throw new Error(`Response status: ${response.status}, status text: ${response.statusText}`);
-			}
-			const data: AutoModeAPIResponse = await response.json() as AutoModeAPIResponse;
-			// HACK: Boost the autoModeHint model to the front of the list until CAPI fixes their bug
-			const hintIndex = data.available_models.indexOf(autoModeHint);
-			if (hintIndex > 0) {
-				data.available_models.splice(hintIndex, 1);
-				data.available_models.unshift(autoModeHint);
-			}
-			this._logService.trace(`Fetched auto model for ${this.debugName} in ${Date.now() - startTime}ms.`);
-			this._token = data;
-			this._usedSinceLastFetch = false;
-			// Trigger a refresh 5 minutes before expiration
-			if (!this._store.isDisposed) {
-				this._refreshTimer.cancelAndSet(() => {
-					if (!this._usedSinceLastFetch) {
-						this._logService.trace(`[${this.debugName}] Skipping auto mode token refresh because it was not used since last fetch.`);
-						this._token = undefined;
-						return;
-					}
-					this._fetchToken();
-				}, (data.expires_at * 1000) - Date.now() - 5 * 60 * 1000);
-			}
-		} catch (err) {
-			this._logService.error(`[${this.debugName}] Failed to fetch AutoMode token:`, err);
-			this._token = undefined;
-		} finally {
-			this._fetchTokenPromise = undefined;
-		}
+		return this._fetchedValue.resolve();
 	}
 }
 

src/platform/endpoint/node/test/automodeService.spec.ts

@@ -21,6 +21,35 @@ import { NullTelemetryService } from '../../../telemetry/common/nullTelemetrySer
 import { ICAPIClientService } from '../../common/capiClient';
 import { AutomodeService } from '../automodeService';
 
+function createMockHeaders(entries: Record<string, string> = {}): { get(name: string): string | null } {
+	const lower: Record<string, string> = {};
+	for (const [k, v] of Object.entries(entries)) {
+		lower[k.toLowerCase()] = v;
+	}
+	return { get: (name: string) => lower[name.toLowerCase()] ?? null };
+}
+
+/**
+ * Creates a mock response with a real stream-backed body so that middleware
+ * cloning (tee) works correctly. Token responses go through the middleware
+ * pipeline where {@link cloneResponse} reads the body stream.
+ */
+function makeMockTokenResponse(body: { available_models: string[]; expires_at: number; session_token: string }) {
+	const serialized = JSON.stringify(body);
+	return {
+		status: 200,
+		headers: createMockHeaders(),
+		body: new ReadableStream<Uint8Array>({
+			start(controller) {
+				controller.enqueue(new TextEncoder().encode(serialized));
+				controller.close();
+			},
+		}),
+		async text() { return serialized; },
+		async json() { return JSON.parse(serialized); },
+	};
+}
+
 describe('AutomodeService', () => {
 	let automodeService: AutomodeService;
 	let mockCAPIClientService: ICAPIClientService;
@@ -64,14 +93,13 @@ describe('AutomodeService', () => {
 	}
 
 	function mockApiResponse(available_models: string[], session_token = 'test-token', expiresInSeconds = 3600): void {
-		(mockCAPIClientService.makeRequest as ReturnType<typeof vi.fn>).mockResolvedValue({
-			ok: true,
-			json: vi.fn().mockResolvedValue({
+		(mockCAPIClientService.makeRequest as ReturnType<typeof vi.fn>).mockResolvedValue(
+			makeMockTokenResponse({
 				available_models,
 				expires_at: Math.floor(Date.now() / 1000) + expiresInSeconds,
 				session_token,
 			})
-		});
+		);
 	}
 
 	function enableRouter(): void {
@@ -85,14 +113,13 @@ describe('AutomodeService', () => {
 		mockChatEndpoint = createEndpoint('gpt-4o-mini', 'OpenAI');
 
 		mockCAPIClientService = {
-			makeRequest: vi.fn().mockResolvedValue({
-				ok: true,
-				json: vi.fn().mockResolvedValue({
+			makeRequest: vi.fn().mockResolvedValue(
+				makeMockTokenResponse({
 					available_models: ['gpt-4o', 'gpt-4o-mini'],
 					expires_at: Math.floor(Date.now() / 1000) + 3600,
 					session_token: 'test-token'
 				})
-			})
+			)
 		} as unknown as ICAPIClientService;
 
 		mockAuthService = {
@@ -154,6 +181,8 @@ describe('AutomodeService', () => {
 				if (opts?.type === RequestType.ModelRouter) {
 					return Promise.resolve({
 						ok: true,
+						status: 200,
+						headers: createMockHeaders(),
 						text: vi.fn().mockResolvedValue(JSON.stringify({
 							predicted_label: 'needs_reasoning',
 							confidence: 0.85,
@@ -165,14 +194,13 @@ describe('AutomodeService', () => {
 						}))
 					});
 				}
-				return Promise.resolve({
-					ok: true,
-					json: vi.fn().mockResolvedValue({
+				return Promise.resolve(
+					makeMockTokenResponse({
 						available_models: ['gpt-4o', 'gpt-4o-mini'],
 						expires_at: Math.floor(Date.now() / 1000) + 3600,
 						session_token: 'test-token'
 					})
-				});
+				);
 			});
 
 			automodeService = createService();
@@ -205,6 +233,8 @@ describe('AutomodeService', () => {
 					capturedBody = req.body;
 					return Promise.resolve({
 						ok: true,
+						status: 200,
+						headers: createMockHeaders(),
 						text: vi.fn().mockResolvedValue(JSON.stringify({
 							predicted_label: 'needs_reasoning',
 							confidence: 0.85,
@@ -216,14 +246,13 @@ describe('AutomodeService', () => {
 						}))
 					});
 				}
-				return Promise.resolve({
-					ok: true,
-					json: vi.fn().mockResolvedValue({
+				return Promise.resolve(
+					makeMockTokenResponse({
 						available_models: ['gpt-4o', 'gpt-4o-mini'],
 						expires_at: Math.floor(Date.now() / 1000) + 3600,
 						session_token: 'test-token'
 					})
-				});
+				);
 			});
 
 			automodeService = createService();
@@ -416,6 +445,8 @@ describe('AutomodeService', () => {
 				if (opts?.type === RequestType.ModelRouter) {
 					return Promise.resolve({
 						ok: true,
+						status: 200,
+						headers: createMockHeaders(),
 						text: vi.fn().mockResolvedValue(JSON.stringify({
 							predicted_label: 'needs_reasoning',
 							confidence: 0.9,
@@ -427,14 +458,13 @@ describe('AutomodeService', () => {
 						}))
 					});
 				}
-				return Promise.resolve({
-					ok: true,
-					json: vi.fn().mockResolvedValue({
+				return Promise.resolve(
+					makeMockTokenResponse({
 						available_models,
 						expires_at: Math.floor(Date.now() / 1000) + 3600,
 						session_token,
 					})
-				});
+				);
 			});
 		}
 
@@ -447,14 +477,13 @@ describe('AutomodeService', () => {
 				if (opts?.type === RequestType.ModelRouter) {
 					return Promise.reject(new Error('Network error'));
 				}
-				return Promise.resolve({
-					ok: true,
-					json: vi.fn().mockResolvedValue({
+				return Promise.resolve(
+					makeMockTokenResponse({
 						available_models: ['claude-sonnet', 'gpt-4o'],
 						expires_at: Math.floor(Date.now() / 1000) + 3600,
 						session_token: 'test-token',
 					})
-				});
+				);
 			});
 
 			automodeService = createService();
@@ -498,14 +527,13 @@ describe('AutomodeService', () => {
 						});
 					});
 				}
-				return Promise.resolve({
-					ok: true,
-					json: vi.fn().mockResolvedValue({
+				return Promise.resolve(
+					makeMockTokenResponse({
 						available_models: ['claude-sonnet', 'gpt-4o'],
 						expires_at: Math.floor(Date.now() / 1000) + 3600,
 						session_token: 'test-token',
 					})
-				});
+				);
 			});
 
 			automodeService = createService();
@@ -669,6 +697,8 @@ describe('AutomodeService', () => {
 				if (opts?.type === RequestType.ModelRouter) {
 					return Promise.resolve({
 						ok: true,
+						status: 200,
+						headers: createMockHeaders(),
 						text: vi.fn().mockResolvedValue(JSON.stringify({
 							predicted_label: 'needs_reasoning',
 							confidence: 0.9,
@@ -680,14 +710,13 @@ describe('AutomodeService', () => {
 						}))
 					});
 				}
-				return Promise.resolve({
-					ok: true,
-					json: vi.fn().mockResolvedValue({
+				return Promise.resolve(
+					makeMockTokenResponse({
 						available_models: ['claude-sonnet', 'gpt-4o'],
 						expires_at: Math.floor(Date.now() / 1000) + 3600,
 						session_token: 'test-token',
 					})
-				});
+				);
 			});
 
 			automodeService = createService();