From a3bdfdaefc91f44e6bb72c7003768bc10b103c64 Mon Sep 17 00:00:00 2001 From: Jackson Chen <90215880+Sma1lboy@users.noreply.github.com> Date: Tue, 5 Nov 2024 02:30:25 -0600 Subject: [PATCH] feat: add global mutexAbortController for chat features (#3372) --- clients/tabby-agent/src/chat/global.ts | 14 ++++++++++++++ clients/tabby-agent/src/chat/inlineEdit.ts | 22 +++++++++++----------- clients/tabby-agent/src/chat/smartApply.ts | 18 +++++++++--------- clients/vscode/src/chat/WebviewHelper.ts | 2 +- 4 files changed, 35 insertions(+), 21 deletions(-) create mode 100644 clients/tabby-agent/src/chat/global.ts diff --git a/clients/tabby-agent/src/chat/global.ts b/clients/tabby-agent/src/chat/global.ts new file mode 100644 index 000000000000..157cbe885fc4 --- /dev/null +++ b/clients/tabby-agent/src/chat/global.ts @@ -0,0 +1,14 @@ +export let mutexAbortController: AbortController | undefined = undefined; + +// reset global mutexAbortController to undefined +export const resetMutexAbortController = () => { + mutexAbortController = undefined; +}; + +// initialize global mutexAbortController +export const initMutexAbortController = () => { + if (!mutexAbortController) { + mutexAbortController = new AbortController(); + } + return mutexAbortController; +}; diff --git a/clients/tabby-agent/src/chat/inlineEdit.ts b/clients/tabby-agent/src/chat/inlineEdit.ts index 3438d76de832..02543641198c 100644 --- a/clients/tabby-agent/src/chat/inlineEdit.ts +++ b/clients/tabby-agent/src/chat/inlineEdit.ts @@ -9,7 +9,6 @@ import { ChatEditRequest, ChatEditParams, ChatEditResolveRequest, - ChatEditResolveParams, ChatEditCommandRequest, ChatEditCommandParams, ChatEditCommand, @@ -17,15 +16,16 @@ import { ChatEditDocumentTooLongError, ChatEditMutexError, ServerCapabilities, + ChatEditResolveParams, } from "../protocol"; import cryptoRandomString from "crypto-random-string"; import { isEmptyRange } from "../utils/range"; -import { applyWorkspaceEdit, readResponseStream, Edit } from "./utils"; +import { readResponseStream, Edit, applyWorkspaceEdit } from "./utils"; +import { initMutexAbortController, mutexAbortController, resetMutexAbortController } from "./global"; export class ChatEditProvider implements Feature { private lspConnection: Connection | undefined = undefined; private currentEdit: Edit | undefined = undefined; - private mutexAbortController: AbortController | undefined = undefined; constructor( private readonly configurations: Configurations, @@ -127,15 +127,15 @@ export class ChatEditProvider implements Feature { throw { name: "ChatEditDocumentTooLongError", message: "Document too long" } as ChatEditDocumentTooLongError; } - if (this.mutexAbortController && !this.mutexAbortController.signal.aborted) { + if (mutexAbortController && !mutexAbortController.signal.aborted) { throw { name: "ChatEditMutexError", message: "Another smart edit is already in progress", } as ChatEditMutexError; } - this.mutexAbortController = new AbortController(); - token.onCancellationRequested(() => this.mutexAbortController?.abort()); + initMutexAbortController(); + token.onCancellationRequested(() => mutexAbortController?.abort()); let insertMode: boolean = isEmptyRange(params.location.range); const presetCommand = /^\/\w+\b/g.exec(params.command)?.[0]; @@ -202,7 +202,7 @@ export class ChatEditProvider implements Feature { model: "", stream: true, }, - this.mutexAbortController.signal, + mutexAbortController?.signal, ); const editId = "tabby-" + cryptoRandomString({ length: 6, type: "alphanumeric" }); @@ -226,10 +226,10 @@ export class ChatEditProvider implements Feature { readableStream, this.lspConnection, this.currentEdit, - this.mutexAbortController, + mutexAbortController, () => { this.currentEdit = undefined; - this.mutexAbortController = undefined; + resetMutexAbortController(); }, config.chat.edit.responseDocumentTag, config.chat.edit.responseCommentTag, @@ -239,13 +239,13 @@ export class ChatEditProvider implements Feature { async stopEdit(id: ChatEditToken): Promise { if (this.isCurrentEdit(id)) { - this.mutexAbortController?.abort(); + mutexAbortController?.abort(); } } async resolveEdit(params: ChatEditResolveParams): Promise { if (params.action === "cancel") { - this.mutexAbortController?.abort(); + mutexAbortController?.abort(); return false; } diff --git a/clients/tabby-agent/src/chat/smartApply.ts b/clients/tabby-agent/src/chat/smartApply.ts index 5947c9255e8c..d983f3064f60 100644 --- a/clients/tabby-agent/src/chat/smartApply.ts +++ b/clients/tabby-agent/src/chat/smartApply.ts @@ -22,12 +22,12 @@ import cryptoRandomString from "crypto-random-string"; import { getLogger } from "../logger"; import { readResponseStream, showDocument, Edit } from "./utils"; import { getSmartApplyRange } from "./smartRange"; +import { initMutexAbortController, mutexAbortController, resetMutexAbortController } from "./global"; const logger = getLogger("SmartApplyFeature"); export class SmartApplyFeature implements Feature { private lspConnection: Connection | undefined = undefined; - private mutexAbortController: AbortController | undefined = undefined; constructor( private readonly configurations: Configurations, private readonly tabbyApiClient: TabbyApiClient, @@ -61,16 +61,16 @@ export class SmartApplyFeature implements Feature { return false; } - if (this.mutexAbortController && !this.mutexAbortController.signal.aborted) { + if (mutexAbortController && !mutexAbortController.signal.aborted) { logger.warn("Another smart edit is already in progress"); throw { name: "ChatEditMutexError", message: "Another smart edit is already in progress", } as ChatEditMutexError; } - this.mutexAbortController = new AbortController(); - logger.debug("mutex abort status: " + (this.mutexAbortController === undefined)); - token.onCancellationRequested(() => this.mutexAbortController?.abort()); + initMutexAbortController(); + logger.debug("mutex abort status: " + (mutexAbortController === undefined)); + token.onCancellationRequested(() => mutexAbortController?.abort()); let applyRange = getSmartApplyRange(document, params.text); //if cannot find range, lets use backend LLMs @@ -115,9 +115,9 @@ export class SmartApplyFeature implements Feature { this.lspConnection, this.tabbyApiClient, this.configurations, - this.mutexAbortController, + mutexAbortController, () => { - this.mutexAbortController = undefined; + resetMutexAbortController(); }, ); return true; @@ -126,7 +126,7 @@ export class SmartApplyFeature implements Feature { return false; } finally { logger.debug("Resetting mutex abort controller"); - this.mutexAbortController = undefined; + resetMutexAbortController(); } } } @@ -222,7 +222,7 @@ async function provideSmartApplyEditLLM( lspConnection: Connection, tabbyApiClient: TabbyApiClient, configurations: Configurations, - mutexAbortController: AbortController, + mutexAbortController: AbortController | undefined, onResetMutex: () => void, ): Promise { if (!document) { diff --git a/clients/vscode/src/chat/WebviewHelper.ts b/clients/vscode/src/chat/WebviewHelper.ts index d8f40580ad82..f87392a0dde6 100644 --- a/clients/vscode/src/chat/WebviewHelper.ts +++ b/clients/vscode/src/chat/WebviewHelper.ts @@ -474,7 +474,7 @@ export class WebviewHelper { this.logger.info("Smart apply in editor started."); this.logger.trace("Smart apply in editor with content:", { content }); - await window.withProgress( + window.withProgress( { location: ProgressLocation.Notification, title: "Smart Apply in Progress",