Better disposable leak fix (microsoft#269613)

Author: lramos15

src/vs/workbench/api/browser/mainThreadLanguageModels.ts

@@ -55,48 +55,53 @@ export class MainThreadLanguageModels implements MainThreadLanguageModelsShape {
 	}
 
 	$registerLanguageModelProvider(vendor: string): void {
-		const disposables = this._store.add(new DisposableStore());
-		disposables.add(this._chatProviderService.registerLanguageModelProvider(vendor, {
-			onDidChange: Event.filter(this._lmProviderChange.event, e => e.vendor === vendor, disposables) as unknown as Event<void>,
-			provideLanguageModelChatInfo: async (options, token) => {
-				const modelsAndIdentifiers = await this._proxy.$provideLanguageModelChatInfo(vendor, options, token);
-				modelsAndIdentifiers.forEach(m => {
-					if (m.metadata.auth) {
-						disposables.add(this._registerAuthenticationProvider(m.metadata.extension, m.metadata.auth));
+		const disposables = new DisposableStore();
+		try {
+			disposables.add(this._chatProviderService.registerLanguageModelProvider(vendor, {
+				onDidChange: Event.filter(this._lmProviderChange.event, e => e.vendor === vendor, disposables) as unknown as Event<void>,
+				provideLanguageModelChatInfo: async (options, token) => {
+					const modelsAndIdentifiers = await this._proxy.$provideLanguageModelChatInfo(vendor, options, token);
+					modelsAndIdentifiers.forEach(m => {
+						if (m.metadata.auth) {
+							disposables.add(this._registerAuthenticationProvider(m.metadata.extension, m.metadata.auth));
+						}
+					});
+					return modelsAndIdentifiers;
+				},
+				sendChatRequest: async (modelId, messages, from, options, token) => {
+					const requestId = (Math.random() * 1e6) | 0;
+					const defer = new DeferredPromise<any>();
+					const stream = new AsyncIterableSource<IChatResponsePart | IChatResponsePart[]>();
+
+					try {
+						this._pendingProgress.set(requestId, { defer, stream });
+						await Promise.all(
+							messages.flatMap(msg => msg.content)
+								.filter(part => part.type === 'image_url')
+								.map(async part => {
+									part.value.data = VSBuffer.wrap(await resizeImage(part.value.data.buffer));
+								})
+						);
+						await this._proxy.$startChatRequest(modelId, requestId, from, new SerializableObjectWithBuffers(messages), options, token);
+					} catch (err) {
+						this._pendingProgress.delete(requestId);
+						throw err;
 					}
-				});
-				return modelsAndIdentifiers;
-			},
-			sendChatRequest: async (modelId, messages, from, options, token) => {
-				const requestId = (Math.random() * 1e6) | 0;
-				const defer = new DeferredPromise<any>();
-				const stream = new AsyncIterableSource<IChatResponsePart | IChatResponsePart[]>();
-
-				try {
-					this._pendingProgress.set(requestId, { defer, stream });
-					await Promise.all(
-						messages.flatMap(msg => msg.content)
-							.filter(part => part.type === 'image_url')
-							.map(async part => {
-								part.value.data = VSBuffer.wrap(await resizeImage(part.value.data.buffer));
-							})
-					);
-					await this._proxy.$startChatRequest(modelId, requestId, from, new SerializableObjectWithBuffers(messages), options, token);
-				} catch (err) {
-					this._pendingProgress.delete(requestId);
-					throw err;
-				}
 
-				return {
-					result: defer.p,
-					stream: stream.asyncIterable
-				} satisfies ILanguageModelChatResponse;
-			},
-			provideTokenCount: (modelId, str, token) => {
-				return this._proxy.$provideTokenLength(modelId, str, token);
-			},
-		}));
-		this._providerRegistrations.set(vendor, disposables);
+					return {
+						result: defer.p,
+						stream: stream.asyncIterable
+					} satisfies ILanguageModelChatResponse;
+				},
+				provideTokenCount: (modelId, str, token) => {
+					return this._proxy.$provideTokenLength(modelId, str, token);
+				},
+			}));
+			this._providerRegistrations.set(vendor, disposables);
+		} catch (err) {
+			disposables.dispose();
+			throw err;
+		}
 	}
 
 	$onLMProviderChange(vendor: string): void {