fix(agent): improve postprocess to trim whitespace. (#550)

* fix(agent): improve postprocess to trim white space.

* fix: lint.

clients/tabby-agent/src/postprocess/index.ts

@@ -4,6 +4,7 @@ import { removeRepetitiveBlocks } from "./removeRepetitiveBlocks";
 import { removeRepetitiveLines } from "./removeRepetitiveLines";
 import { removeLineEndsWithRepetition } from "./removeLineEndsWithRepetition";
 import { limitScopeByIndentation } from "./limitScopeByIndentation";
+import { trimSpace } from "./trimSpace";
 import { removeOverlapping } from "./removeOverlapping";
 import { dropDuplicated } from "./dropDuplicated";
 import { dropBlank } from "./dropBlank";
@@ -15,8 +16,9 @@ export async function preCacheProcess(
   const context = buildContext(request);
   return Promise.resolve(response)
     .then(applyFilter(removeLineEndsWithRepetition(context)))
-    .then(applyFilter(removeOverlapping(context)))
     .then(applyFilter(dropDuplicated(context)))
+    .then(applyFilter(trimSpace(context)))
+    .then(applyFilter(removeOverlapping(context)))
     .then(applyFilter(dropBlank()));
 }
 
@@ -29,5 +31,7 @@ export async function postprocess(
     .then(applyFilter(removeRepetitiveBlocks(context)))
     .then(applyFilter(removeRepetitiveLines(context)))
     .then(applyFilter(limitScopeByIndentation(context)))
+    .then(applyFilter(trimSpace(context)))
+    .then(applyFilter(removeOverlapping(context)))
     .then(applyFilter(dropBlank()));
 }

clients/tabby-agent/src/postprocess/trimSpace.test.ts

@@ -0,0 +1,83 @@
+import { expect } from "chai";
+import { documentContext, inline } from "./testUtils";
+import { trimSpace } from "./trimSpace";
+
+describe("postprocess", () => {
+  describe("trimSpace", () => {
+    it("should remove trailing space", () => {
+      const context = {
+        ...documentContext`
+        let foo = new ║
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+                      ├Foo(); ┤
+        `;
+      const expected = inline`
+                      ├Foo();┤
+      `;
+      expect(trimSpace(context)(completion)).to.eq(expected);
+    });
+
+    it("should not remove trailing space if filling in line", () => {
+      const context = {
+        ...documentContext`
+        let foo = sum(║baz)
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+                      ├bar, ┤
+        `;
+      expect(trimSpace(context)(completion)).to.eq(completion);
+    });
+
+    it("should remove trailing space if filling in line with suffix starts with space", () => {
+      const context = {
+        ...documentContext`
+        let foo = sum(║ baz)
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+                      ├bar, ┤
+        `;
+      const expected = inline`
+                      ├bar,┤
+        `;
+      expect(trimSpace(context)(completion)).to.eq(expected);
+    });
+
+    it("should not remove leading space if current line is blank", () => {
+      const context = {
+        ...documentContext`
+        function sum(a, b) {
+        ║
+        }
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+        ├  return a + b;┤
+        `;
+      expect(trimSpace(context)(completion)).to.eq(completion);
+    });
+
+    it("should remove leading space if current line is not blank and ends with space", () => {
+      const context = {
+        ...documentContext`
+        let foo = ║
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+                  ├ sum(bar, baz);┤
+        `;
+      const expected = inline`
+                  ├sum(bar, baz);┤
+        `;
+      expect(trimSpace(context)(completion)).to.eq(expected);
+    });
+  });
+});

clients/tabby-agent/src/postprocess/trimSpace.ts

@@ -0,0 +1,24 @@
+import { PostprocessFilter, PostprocessContext } from "./base";
+import { splitLines, isBlank } from "../utils";
+
+export const trimSpace: (context: PostprocessContext) => PostprocessFilter = (context) => {
+  return (input) => {
+    const { prefixLines, suffixLines } = context;
+    const inputLines = splitLines(input);
+    let trimmedInput = input;
+    const prefixCurrentLine = prefixLines[prefixLines.length - 1] ?? "";
+    const suffixCurrentLine = suffixLines[0] ?? "";
+    if (!isBlank(prefixCurrentLine) && prefixCurrentLine.match(/\s$/)) {
+      trimmedInput = trimmedInput.trimStart();
+    }
+
+    if (
+      inputLines.length > 1 ||
+      isBlank(suffixCurrentLine) ||
+      (!isBlank(suffixCurrentLine) && suffixCurrentLine.match(/^\s/))
+    ) {
+      trimmedInput = trimmedInput.trimEnd();
+    }
+    return trimmedInput;
+  };
+};

clients/tabby-agent/src/utils.ts

@@ -1,5 +1,10 @@
 export function splitLines(input: string) {
-  return input.match(/.*(?:$|\r?\n)/g).filter(Boolean); // Split lines and keep newline character
+  const lines = input.match(/.*(?:$|\r?\n)/g).filter(Boolean); // Split lines and keep newline character
+  if (lines.length > 0 && lines[lines.length - 1].endsWith("\n")) {
+    // Keep last empty line
+    lines.push("");
+  }
+  return lines;
 }
 
 export function splitWords(input: string) {