add ability to abort

src/vs/platform/void/browser/llmMessageService.ts

@@ -3,13 +3,14 @@
  *  Void Editor additions licensed under the AGPLv3 License.
  *--------------------------------------------------------------------------------------------*/
 
-import { ProxyOnTextPayload, ProxyOnErrorPayload, ProxyOnFinalMessagePayload, LLMMessageServiceParams, ProxyLLMMessageParams } from '../common/llmMessageTypes.js';
+import { ProxyOnTextPayload, ProxyOnErrorPayload, ProxyOnFinalMessagePayload, LLMMessageServiceParams, ProxyLLMMessageParams, ProxyLLMMessageAbortParams } from '../common/llmMessageTypes.js';
 import { IChannel } from '../../../base/parts/ipc/common/ipc.js';
 import { IMainProcessService } from '../../ipc/common/mainProcessService.js';
 import { InstantiationType, registerSingleton } from '../../instantiation/common/extensions.js';
 import { generateUuid } from '../../../base/common/uuid.js';
 import { createDecorator } from '../../instantiation/common/instantiation.js';
 import { Event } from '../../../base/common/event.js';
+import { IDisposable } from '../../../base/common/lifecycle.js';
 
 
 // BROWSER IMPLEMENTATION OF SENDLLMMESSAGE
@@ -23,10 +24,11 @@ export interface ISendLLMMessageService {
 
 
 export class SendLLMMessageService implements ISendLLMMessageService {
-	static readonly ID = 'void.contrib.browserSendLLMMessageService';
 
 	readonly _serviceBrand: undefined;
-	readonly channel: IChannel;
+	private readonly channel: IChannel;
+
+	private readonly _disposablesOfRequestId: Record<string, IDisposable[]> = {}
 
 	constructor(
 		@IMainProcessService mainProcessService: IMainProcessService // used as a renderer (only usable on client side)
@@ -36,6 +38,15 @@ export class SendLLMMessageService implements ISendLLMMessageService {
 		// const service = ProxyChannel.toService<LLMMessageChannel>(mainProcessService.getChannel('void-channel-sendLLMMessage')); // lets you call it like a service, not needed here
 	}
 
+	_addDisposable(requestId: string, disposable: IDisposable) {
+		if (!this._disposablesOfRequestId[requestId]) {
+			this._disposablesOfRequestId[requestId] = []
+		}
+		this._disposablesOfRequestId[requestId].push(disposable)
+	}
+
+
+
 	sendLLMMessage(params: LLMMessageServiceParams) {
 		const requestId_ = generateUuid();
 		const { onText, onFinalMessage, onError, ...proxyParams } = params;
@@ -43,25 +54,48 @@ export class SendLLMMessageService implements ISendLLMMessageService {
 		// listen for listenerName='onText' | 'onFinalMessage' | 'onError', and call the original function on it
 
 		const onTextEvent: Event<ProxyOnTextPayload> = this.channel.listen('onText')
-		onTextEvent(e => {
-			if (requestId_ !== e.requestId) return;
-			onText(e)
-		})
+		this._addDisposable(requestId_,
+			onTextEvent(e => {
+				if (requestId_ !== e.requestId) return;
+				onText(e)
+			})
+		)
 
 		const onFinalMessageEvent: Event<ProxyOnFinalMessagePayload> = this.channel.listen('onFinalMessage')
-		onFinalMessageEvent(e => {
-			if (requestId_ !== e.requestId) return;
-			onFinalMessage(e)
-		})
+		this._addDisposable(requestId_,
+			onFinalMessageEvent(e => {
+				if (requestId_ !== e.requestId) return;
+				onFinalMessage(e)
+				this._dispose(requestId_)
+			})
+		)
 
 		const onErrorEvent: Event<ProxyOnErrorPayload> = this.channel.listen('onError')
-		onErrorEvent(e => {
-			if (requestId_ !== e.requestId) return;
-			onError(e)
-		})
+		this._addDisposable(requestId_,
+			onErrorEvent(e => {
+				if (requestId_ !== e.requestId) return;
+				onError(e)
+				this._dispose(requestId_)
+			})
+		)
 
 		// params will be stripped of all its functions
 		this.channel.call('sendLLMMessage', { ...proxyParams, requestId: requestId_ } satisfies ProxyLLMMessageParams);
+
+		return requestId_
+	}
+
+	private _dispose(requestId: string) {
+		if (!(requestId in this._disposablesOfRequestId)) return
+		for (const disposable of this._disposablesOfRequestId[requestId]) {
+			disposable.dispose()
+		}
+		delete this._disposablesOfRequestId[requestId]
+	}
+
+	abort(requestId: string) {
+		this.channel.call('abort', { requestId } satisfies ProxyLLMMessageAbortParams);
+		this._dispose(requestId)
 	}
 }
 

src/vs/platform/void/common/llmMessageTypes.ts

@@ -27,13 +27,26 @@ export type LLMMessageServiceParams = {
 
 	messages: LLMMessage[];
 	voidConfig: VoidConfig | null;
-	abortRef: AbortRef;
 
 	logging: {
 		loggingName: string,
 	};
 }
 
+export type SendLLMMMessageParams = {
+	onText: OnText;
+	onFinalMessage: OnFinalMessage;
+	onError: OnError;
+
+	messages: LLMMessage[];
+	voidConfig: VoidConfig | null;
+
+	logging: {
+		loggingName: string,
+	};
+	abortRef: AbortRef;
+}
+
 // can't send functions across a proxy, use listeners instead
 export const listenerNames = ['onText', 'onFinalMessage', 'onError'] as const
 export type ProxyLLMMessageParams = Omit<LLMMessageServiceParams, typeof listenerNames[number]> & { requestId: string }
@@ -41,3 +54,5 @@ export type ProxyLLMMessageParams = Omit<LLMMessageServiceParams, typeof listene
 export type ProxyOnTextPayload = Parameters<OnText>[0] & { requestId: string }
 export type ProxyOnFinalMessagePayload = Parameters<OnFinalMessage>[0] & { requestId: string }
 export type ProxyOnErrorPayload = Parameters<OnError>[0] & { requestId: string }
+
+export type ProxyLLMMessageAbortParams = { requestId: string }

src/vs/platform/void/electron-main/LLMMessageChannel.ts

@@ -11,7 +11,7 @@
 import { IServerChannel } from '../../../base/parts/ipc/common/ipc.js';
 import { Emitter, Event } from '../../../base/common/event.js';
 import { sendLLMMessage } from '../../../workbench/contrib/void/browser/react/out/util/sendLLMMessage.js';
-import { listenerNames, ProxyOnTextPayload, ProxyOnErrorPayload, ProxyOnFinalMessagePayload, LLMMessageServiceParams, ProxyLLMMessageParams } from '../common/llmMessageTypes.js';
+import { listenerNames, ProxyOnTextPayload, ProxyOnErrorPayload, ProxyOnFinalMessagePayload, ProxyLLMMessageParams, AbortRef, SendLLMMMessageParams, ProxyLLMMessageAbortParams } from '../common/llmMessageTypes.js';
 
 // NODE IMPLEMENTATION OF SENDLLMMESSAGE - calls sendLLMMessage() and returns listeners
 
@@ -25,11 +25,14 @@ export class LLMMessageChannel implements IServerChannel {
 	private readonly _onError = new Emitter<ProxyOnErrorPayload>();
 	readonly onError = this._onError.event;
 
+
+	private readonly _abortRefOfRequestId: Record<string, AbortRef> = {}
+
+
 	constructor() { }
 
-	// browser uses this
+	// browser uses this to listen for changes
 	listen(_: unknown, event: typeof listenerNames[number]): Event<any> {
-		console.log('event LISTENING!!!:', event)
 		if (event === 'onText') {
 			return this.onText;
 		}
@@ -44,23 +47,48 @@ export class LLMMessageChannel implements IServerChannel {
 		}
 	}
 
-	// both use this
-	async call(_: unknown, command: string, params: ProxyLLMMessageParams): Promise<any> {
-
-		if (command !== 'sendLLMMessage') throw new Error(`Invalid call in sendLLMMessage channel: ${command}.\nArgs:\n${JSON.stringify(params, null, 5)}`);
+	// browser uses this to call
+	async call(_: unknown, command: string, params: any): Promise<any> {
 
 		try {
-			const { requestId } = params;
-			const mainThreadParams: LLMMessageServiceParams = {
-				...params,
-				onText: ({ newText, fullText }) => { this._onText.fire({ requestId, newText, fullText }); },
-				onFinalMessage: ({ fullText }) => { this._onFinalMessage.fire({ requestId, fullText }); },
-				onError: ({ error }) => { this._onError.fire({ requestId, error }); },
+			if (command === 'sendLLMMessage') {
+				this._callSendLLMMessage(params)
+			}
+			else if (command === 'abort') {
+				this._callAbort(params)
+			}
+			else {
+				throw new Error(`Void sendLLM: command "${command}" not recognized.`)
 			}
-			sendLLMMessage(mainThreadParams);
 		}
 		catch (e) {
 			console.log('sendLLM channel: call error', e)
 		}
 	}
+
+	private _callSendLLMMessage(params: ProxyLLMMessageParams) {
+		const { requestId } = params;
+
+		if (!(requestId in this._abortRefOfRequestId))
+			this._abortRefOfRequestId[requestId] = { current: null }
+
+		const mainThreadParams: SendLLMMMessageParams = {
+			...params,
+			onText: ({ newText, fullText }) => { this._onText.fire({ requestId, newText, fullText }); },
+			onFinalMessage: ({ fullText }) => { this._onFinalMessage.fire({ requestId, fullText }); },
+			onError: ({ error }) => { this._onError.fire({ requestId, error }); },
+			abortRef: this._abortRefOfRequestId[requestId],
+		}
+		sendLLMMessage(mainThreadParams);
+	}
+
+
+	private _callAbort(params: ProxyLLMMessageAbortParams) {
+		const { requestId } = params;
+		if (!(requestId in this._abortRefOfRequestId)) return
+		this._abortRefOfRequestId[requestId].current?.()
+		delete this._abortRefOfRequestId[requestId]
+	}
+
+
 }

src/vs/workbench/contrib/void/browser/react/src/util/sendLLMMessage.tsx

@@ -4,7 +4,7 @@ import { Ollama } from 'ollama/browser'
 import { Content, GoogleGenerativeAI, GoogleGenerativeAIFetchError } from '@google/generative-ai';
 import { posthog } from 'posthog-js'
 import type { VoidConfig } from '../../../registerConfig.js';
-import type { LLMMessage, OnText, OnError, OnFinalMessage, } from '../../../../../../../platform/void/common/llmMessageTypes.js';
+import type { LLMMessage, OnText, OnError, OnFinalMessage, SendLLMMMessageParams, } from '../../../../../../../platform/void/common/llmMessageTypes.js';
 import { LLMMessageServiceParams } from '../../../../../../../platform/void/common/llmMessageTypes.js';
 
 type SendLLMMessageFnTypeInternal = (params: {
@@ -284,7 +284,7 @@ export const sendLLMMessage = ({
 	abortRef: abortRef_,
 	voidConfig,
 	logging: { loggingName }
-}: LLMMessageServiceParams) => {
+}: SendLLMMMessageParams) => {
 	if (!voidConfig) return;
 
 	// trim message content (Anthropic and other providers give an error if there is trailing whitespace)

src/vs/workbench/contrib/void/browser/registerInlineDiffs.ts

@@ -761,7 +761,6 @@ Please finish writing the new file by applying the diff to the original file. Re
 					resolve();
 				},
 				voidConfig,
-				abortRef,
 			}
 
 			this._sendLLMMessageService.sendLLMMessage(object)