Skip to content

Commit

Permalink
feat(vscode): one undo-redo step for one chat stream (#2872)
Browse files Browse the repository at this point in the history 
* feat(tabby-agent): adding tabby/workspace/applyEdi tsserver to client namespace

* feat:  send request by using new tabby/workspace/applyEdit type in ChatEditProvider to applyEdit

* feat(vscode):  onRequest by using new tabby/workspace/applyEdit type in ChatFeature to applyEdit

* feat(client): adding middleware to prevent that clients not support specific connection type

* to(vscode): remove inner applyUndo func and adding optional value check

* to(client): remove console log and fix cannot decode GENERATEDCODE tag
  • Loading branch information
@Sma1lboy
Sma1lboy authored Aug 26, 2024
1 parent 17d3b4a commit ecc5bfb
Show file tree
Hide file tree
Showing 3 changed files with 199 additions and 99 deletions.
210 changes: 113 additions & 97 deletions clients/tabby-agent/src/lsp/ChatEditProvider.ts
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { Range, Location, Connection, CancellationToken } from "vscode-languageserver";
import { Range, Location, Connection, CancellationToken, WorkspaceEdit } from "vscode-languageserver";
import {
ChatEditToken,
ChatEditRequest,
Expand All @@ -12,6 +12,8 @@ import {
ChatEditDocumentTooLongError,
ChatEditCommandTooLongError,
ChatEditMutexError,
ApplyWorkspaceEditRequest,
ApplyWorkspaceEditParams,
} from "./protocol";
import { TextDocuments } from "./TextDocuments";
import { TextDocument } from "vscode-languageserver-textdocument";
Expand Down Expand Up @@ -229,7 +231,8 @@ export class ChatEditProvider {
}
}
});
await this.connection.workspace.applyEdit({

await this.applyWorkspaceEdit({
edit: {
changes: {
[params.location.uri]: [
Expand All @@ -240,6 +243,10 @@ export class ChatEditProvider {
],
},
},
options: {
undoStopBefore: false,
undoStopAfter: false,
},
});
return true;
}
Expand All @@ -249,121 +256,130 @@ export class ChatEditProvider {
responseDocumentTag: string[],
responseCommentTag?: string[],
): Promise<void> {
const finalize = async (state: "completed" | "stopped") => {
if (this.currentEdit) {
const edit = this.currentEdit;
edit.state = state;
const editedLines = this.generateChangesPreview(edit);
await this.connection.workspace.applyEdit({
edit: {
changes: {
[edit.location.uri]: [
{
range: edit.editedRange,
newText: editedLines.join("\n") + "\n",
},
],
const applyEdit = async (edit: Edit, isFirst: boolean = false, isLast: boolean = false) => {
const editedLines = this.generateChangesPreview(edit);
const workspaceEdit: WorkspaceEdit = {
changes: {
[edit.location.uri]: [
{
range: edit.editedRange,
newText: editedLines.join("\n") + "\n",
},
},
});
],
},
};

await this.applyWorkspaceEdit({
edit: workspaceEdit,
options: {
undoStopBefore: isFirst,
undoStopAfter: isLast,
},
});

edit.editedRange = {
start: { line: edit.editedRange.start.line, character: 0 },
end: { line: edit.editedRange.start.line + editedLines.length, character: 0 },
};
};

const processBuffer = (edit: Edit, inTag: "document" | "comment", openTag: string, closeTag: string) => {
if (edit.buffer.startsWith(openTag)) {
edit.buffer = edit.buffer.substring(openTag.length);
}
this.currentEdit = null;
this.mutexAbortController = null;

const reg = this.createCloseTagMatcher(closeTag);
const match = reg.exec(edit.buffer);
if (!match) {
edit[inTag === "document" ? "editedText" : "comments"] += edit.buffer;
edit.buffer = "";
} else {
edit[inTag === "document" ? "editedText" : "comments"] += edit.buffer.substring(0, match.index);
edit.buffer = edit.buffer.substring(match.index);
return match[0] === closeTag ? false : inTag;
}
return inTag;
};
const findOpenTag = (
buffer: string,
responseDocumentTag: string[],
responseCommentTag?: string[],
): "document" | "comment" | false => {
const openTags = [responseDocumentTag[0], responseCommentTag?.[0]].filter(Boolean);
if (openTags.length < 1) return false;

const reg = new RegExp(openTags.join("|"), "g");
const match = reg.exec(buffer);
if (match && match[0]) {
if (match[0] === responseDocumentTag[0]) {
return "document";
} else if (match[0] === responseCommentTag?.[0]) {
return "comment";
}
}
return false;
};

try {
let inTag: "document" | "comment" | false = false;
let isFirstEdit = true;

for await (const delta of stream) {
if (!this.currentEdit || !this.mutexAbortController || this.mutexAbortController.signal.aborted) {
break;
}
let changed = false;

const edit = this.currentEdit;
edit.buffer += delta;

if (!inTag) {
const openTags = [responseDocumentTag[0], responseCommentTag?.[0]].filter(Boolean);
if (openTags.length < 1) {
break;
}
const reg = new RegExp(openTags.join("|"), "g");
const match = reg.exec(edit.buffer);
if (match && match[0]) {
if (match[0] === responseDocumentTag[0]) {
inTag = "document";
edit.buffer = edit.buffer.substring(match.index + match[0].length);
} else if (match[0] === responseCommentTag?.[0]) {
inTag = "comment";
edit.buffer = edit.buffer.substring(match.index + match[0].length);
}
}
inTag = findOpenTag(edit.buffer, responseDocumentTag, responseCommentTag);
}

if (inTag) {
let closeTag: string | undefined = undefined;
if (inTag === "document") {
closeTag = responseDocumentTag[1];
} else if (inTag === "comment") {
closeTag = responseCommentTag?.[1];
}
if (!closeTag) {
break;
}
const reg = this.createCloseTagMatcher(closeTag);
const match = reg.exec(edit.buffer);
if (!match) {
if (inTag === "document") {
edit.editedText += edit.buffer;
} else if (inTag === "comment") {
edit.comments += edit.buffer;
}
edit.buffer = "";
} else {
if (inTag === "document") {
edit.editedText += edit.buffer.substring(0, match.index);
} else if (inTag === "comment") {
edit.comments += edit.buffer.substring(0, match.index);
}
edit.buffer = edit.buffer.substring(match.index);
if (match[0] === closeTag) {
inTag = false;
}
}
changed = true;
}
if (changed) {
const editedLines = this.generateChangesPreview(edit);
await this.connection.workspace.applyEdit({
edit: {
changes: {
[edit.location.uri]: [
{
range: edit.editedRange,
newText: editedLines.join("\n") + "\n",
},
],
},
},
});
edit.editedRange = {
start: {
line: edit.editedRange.start.line,
character: 0,
},
end: {
line: edit.editedRange.start.line + editedLines.length,
character: 0,
},
};
const openTag = inTag === "document" ? responseDocumentTag[0] : responseCommentTag?.[0];
const closeTag = inTag === "document" ? responseDocumentTag[1] : responseCommentTag?.[1];
if (!closeTag || !openTag) break;
inTag = processBuffer(edit, inTag, openTag, closeTag);
await applyEdit(edit, isFirstEdit, false);
isFirstEdit = false;
}
}

if (this.currentEdit) {
this.currentEdit.state = "completed";
await applyEdit(this.currentEdit, false, true);
}
} catch (error) {
await finalize("stopped");
// FIXME(@icycodes): use openai for nodejs instead of tabby-openapi schema
if (error instanceof TypeError && error.message.startsWith("terminated")) {
// ignore server side close error
} else {
if (this.currentEdit) {
this.currentEdit.state = "stopped";
await applyEdit(this.currentEdit, false, true);
}
if (!(error instanceof TypeError && error.message.startsWith("terminated"))) {
throw error;
}
} finally {
this.currentEdit = null;
this.mutexAbortController = null;
}
}

private async applyWorkspaceEdit(params: ApplyWorkspaceEditParams): Promise<boolean> {
try {
//TODO(Sma1lboy): adding client capabilities to indicate if client support this method rather than try-catch
const result = await this.connection.sendRequest(ApplyWorkspaceEditRequest.type, params);
return result;
} catch (error) {
try {
await this.connection.workspace.applyEdit({
edit: params.edit,
label: params.label,
});
return true;
} catch (fallbackError) {
return false;
}
}
await finalize("completed");
}

// header line
Expand Down
38 changes: 38 additions & 0 deletions clients/tabby-agent/src/lsp/protocol.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ import {
SemanticTokensRangeParams,
SemanticTokens,
SemanticTokensLegend,
WorkspaceEdit,
} from "vscode-languageserver-protocol";

/**
Expand Down Expand Up @@ -471,6 +472,43 @@ export type ChatEditResolveParams = {
action: "accept" | "discard";
};

/**
* [Tabby] Apply workspace edit request(↩️)
*
* This method is sent from the server to client to apply edit in workspace with options.
* - method: `tabby/workspace/applyEdit`
* - params: {@link ApplyWorkspaceEditParams}
* - result: boolean
*/
export namespace ApplyWorkspaceEditRequest {
export const method = "tabby/workspace/applyEdit";
export const messageDirection = MessageDirection.serverToClient;
export const type = new ProtocolRequestType<ApplyWorkspaceEditParams, boolean, never, void, void>(method);
}

export interface ApplyWorkspaceEditParams {
/**
* An optional label of the workspace edit. This label is
* presented in the user interface for example on an undo
* stack to undo the workspace edit.
*/
label?: string;
/**
* The edits to apply.
*/
edit: WorkspaceEdit;
options?: {
/**
* Add undo stop before making the edits.
*/
readonly undoStopBefore: boolean;
/**
* Add undo stop after making the edits.
*/
readonly undoStopAfter: boolean;
};
}

export type ChatEditResolveCommand = LspCommand & {
title: string;
tooltip?: string;
Expand Down
Loading

0 comments on commit ecc5bfb

Please sign in to comment.