diff --git a/clients/tabby-agent/src/chat/SmartApply.ts b/clients/tabby-agent/src/chat/SmartApply.ts index 6ed07abea6d2..7f7fdea7c8ce 100644 --- a/clients/tabby-agent/src/chat/SmartApply.ts +++ b/clients/tabby-agent/src/chat/SmartApply.ts @@ -24,13 +24,15 @@ import { Configurations } from "../config"; import { TabbyApiClient } from "../http/tabbyApiClient"; import cryptoRandomString from "crypto-random-string"; import { getLogger } from "../logger"; -import { ChatStatus } from "./chatStatus"; import { applyWorkspaceEdit, readResponseStream, revealEditorRange } from "./utils"; import { TextDocument } from "vscode-languageserver-textdocument"; import { getSmartApplyRange } from "./SmartRange"; +import { Edit } from "./inlineEdit"; export class SmartApplyFeature implements Feature { private logger = getLogger("ChatEditProvider"); private lspConnection: Connection | undefined = undefined; + private currentEdit: Edit | undefined = undefined; + private mutexAbortController: AbortController | undefined = undefined; constructor( private readonly configurations: Configurations, private readonly tabbyApiClient: TabbyApiClient, @@ -52,7 +54,7 @@ export class SmartApplyFeature implements Feature { //nothing } - async provideSmartApplyEdit(params: SmartApplyCodeParams, _token: CancellationToken): Promise { + async provideSmartApplyEdit(params: SmartApplyCodeParams, token: CancellationToken): Promise { this.logger.info("Getting document"); const document = this.documents.get(params.location.uri); if (!document) { @@ -64,13 +66,15 @@ export class SmartApplyFeature implements Feature { return false; } - if (ChatStatus.mutexAbortController && !ChatStatus.mutexAbortController.signal.aborted) { + if (this.mutexAbortController && !this.mutexAbortController.signal.aborted) { this.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(); + token.onCancellationRequested(() => this.mutexAbortController?.abort()); let applyRange = getSmartApplyRange(document, params.applyCode); //if cannot find range, lets use backend LLMs @@ -124,7 +128,7 @@ export class SmartApplyFeature implements Feature { edit: workspaceEdit, }; - const revealEditorRangeParams : RevealEditorRangeParams = { + const revealEditorRangeParams: RevealEditorRangeParams = { range: edit.range, revealType: TextEditorRevealType.InCenterIfOutsideViewport, }; @@ -139,7 +143,7 @@ export class SmartApplyFeature implements Feature { return false; } finally { this.logger.info("Resetting mutex abort controller"); - ChatStatus.mutexAbortController = undefined; + this.mutexAbortController = undefined; } } @@ -165,13 +169,6 @@ export class SmartApplyFeature implements Feature { .map((line, idx) => `${idx + 1} | ${line}`) .join("\n"); - if (ChatStatus.mutexAbortController && !ChatStatus.mutexAbortController.signal.aborted) { - throw { - name: "ChatEditMutexError", - message: "Another edit is already in progress", - } as ChatEditMutexError; - } - const config = this.configurations.getMergedConfig(); const promptTemplate = config.chat.provideSmartApplyLineRange.promptTemplate; @@ -265,6 +262,15 @@ export class SmartApplyFeature implements Feature { throw { name: "ChatEditDocumentTooLongError", message: "Document too long" } as ChatEditDocumentTooLongError; } + if (this.mutexAbortController && !this.mutexAbortController.signal.aborted) { + this.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(); + const insertMode = location.range.start.line === location.range.end.line; const presetConfig = config.chat.edit.presetCommands["/smartApply"]; @@ -323,7 +329,7 @@ export class SmartApplyFeature implements Feature { } const editId = "tabby-" + cryptoRandomString({ length: 6, type: "alphanumeric" }); - ChatStatus.currentEdit = { + this.currentEdit = { id: editId, location: location, languageId: document.languageId, @@ -340,6 +346,12 @@ export class SmartApplyFeature implements Feature { await readResponseStream( readableStream, this.lspConnection, + this.currentEdit, + this.mutexAbortController, + () => { + this.currentEdit = undefined; + this.mutexAbortController = undefined; + }, config.chat.edit.responseDocumentTag, config.chat.edit.responseCommentTag, ); @@ -347,6 +359,8 @@ export class SmartApplyFeature implements Feature { return true; } catch (error) { return false; + } finally { + this.mutexAbortController = undefined; } } } diff --git a/clients/tabby-agent/src/chat/chatStatus.ts b/clients/tabby-agent/src/chat/chatStatus.ts deleted file mode 100644 index 0c27165de5e9..000000000000 --- a/clients/tabby-agent/src/chat/chatStatus.ts +++ /dev/null @@ -1,6 +0,0 @@ -import { Edit } from "./inlineEdit"; - -export class ChatStatus { - public static currentEdit: Edit | undefined = undefined; - public static mutexAbortController: AbortController | undefined = undefined; -} diff --git a/clients/tabby-agent/src/chat/inlineEdit.ts b/clients/tabby-agent/src/chat/inlineEdit.ts index 62016a8943ba..5194a0b21830 100644 --- a/clients/tabby-agent/src/chat/inlineEdit.ts +++ b/clients/tabby-agent/src/chat/inlineEdit.ts @@ -237,6 +237,12 @@ export class ChatEditProvider implements Feature { await readResponseStream( readableStream, this.lspConnection, + this.currentEdit, + this.mutexAbortController, + () => { + this.currentEdit = undefined; + this.mutexAbortController = undefined; + }, config.chat.edit.responseDocumentTag, config.chat.edit.responseCommentTag, ); diff --git a/clients/tabby-agent/src/chat/utils.ts b/clients/tabby-agent/src/chat/utils.ts index 1b944bd39b38..868a17f3282f 100644 --- a/clients/tabby-agent/src/chat/utils.ts +++ b/clients/tabby-agent/src/chat/utils.ts @@ -12,11 +12,13 @@ import { RevealEditorRangeParams, RevealEditorRangeRequest, } from "../protocol"; -import { ChatStatus } from "./chatStatus"; export async function readResponseStream( stream: Readable, connection: Connection, + currentEdit: Edit | undefined, + mutexAbortController: AbortController | undefined, + resetEditAndMutexAbortController: () => void, responseDocumentTag: string[], responseCommentTag?: string[], ): Promise { @@ -120,21 +122,21 @@ export async function readResponseStream( }; try { - if (!ChatStatus.currentEdit) { + if (!currentEdit) { throw new Error("No current edit"); } let inTag: "document" | "comment" | false = false; // Insert the first line as early as possible so codelens can be shown - await applyEdit(ChatStatus.currentEdit, true, false); + await applyEdit(currentEdit, true, false); for await (const item of stream) { - if (!ChatStatus.mutexAbortController || ChatStatus.mutexAbortController.signal.aborted) { + if (!mutexAbortController || mutexAbortController.signal.aborted) { break; } const delta = typeof item === "string" ? item : ""; - const edit = ChatStatus.currentEdit; + const edit = currentEdit; edit.buffer += delta; if (!inTag) { @@ -152,21 +154,20 @@ export async function readResponseStream( } } - if (ChatStatus.currentEdit) { - ChatStatus.currentEdit.state = "completed"; - await applyEdit(ChatStatus.currentEdit, false, true); + if (currentEdit) { + currentEdit.state = "completed"; + await applyEdit(currentEdit, false, true); } } catch (error) { - if (ChatStatus.currentEdit) { - ChatStatus.currentEdit.state = "stopped"; - await applyEdit(ChatStatus.currentEdit, false, true); + if (currentEdit) { + currentEdit.state = "stopped"; + await applyEdit(currentEdit, false, true); } if (!(error instanceof TypeError && error.message.startsWith("terminated"))) { throw error; } } finally { - ChatStatus.currentEdit = undefined; - ChatStatus.mutexAbortController = undefined; + resetEditAndMutexAbortController(); } } @@ -198,6 +199,7 @@ export async function revealEditorRange(params: RevealEditorRangeParams, lspConn if (!lspConnection) { return false; } + try { const result = await lspConnection.sendRequest(RevealEditorRangeRequest.type, params); return result;