fix(agent): improve postprocess: limitScopeByIndentaion. (#367)

* fix(agent): improve postprocess: limitScopeByIndentaion.

* fix lint.

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

@@ -4,7 +4,74 @@ import { limitScopeByIndentation } from "./limitScopeByIndentation";
 
 describe("postprocess", () => {
   describe("limitScopeByIndentation", () => {
-    it("should remove content out of current intent scope", () => {
+    it("should drop multiline completions, when the suffix have meaningful chars in the current line.", () => {
+      const context = {
+        ...documentContext`
+        let error = new Error("Something went wrong");
+        console.log(║message);
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+                    ├message);
+        throw error;┤
+      `;
+      expect(limitScopeByIndentation(context)(completion)).to.be.null;
+    });
+
+    it("should allow singleline completions, when the suffix have meaningful chars in the current line.", () => {
+      const context = {
+        ...documentContext`
+        let error = new Error("Something went wrong");
+        console.log(║message);
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+                    ├error, ┤
+      `;
+      expect(limitScopeByIndentation(context)(completion)).to.eq(completion);
+    });
+
+    it("should allow multiline completions, when the suffix only have special chars that will be replaced in the current line, such as `)]}`.", () => {
+      const context = {
+        ...documentContext`
+        function findMax(arr) {║}
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+                               ├
+          let max = arr[0];
+          for (let i = 1; i < arr.length; i++) {
+            if (arr[i] > max) {
+              max = arr[i];
+            }
+          }
+          return max;
+        }┤
+      `;
+      expect(limitScopeByIndentation(context)(completion)).to.eq(completion);
+    });
+
+    it("should limit scope at sentence end, when completion is continuing uncompleted sentence in the prefix.", () => {
+      const context = {
+        ...documentContext`
+        let a =║
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+               ├ 1;
+        let b = 2;┤
+      `;
+      const expected = inline`
+               ├ 1;┤
+      `;
+      expect(limitScopeByIndentation(context)(completion)).to.eq(expected);
+    });
+
+    it("should limit scope at sentence end, when completion is continuing uncompleted sentence in the prefix.", () => {
       const context = {
         ...documentContext`
         function safeParse(json) {
@@ -27,14 +94,108 @@ describe("postprocess", () => {
         }┤
       `;
       const expected = inline`
-                        ├("Parsing", { json });
+                        ├("Parsing", { json });┤
-            return JSON.parse(json);┤
-        ┴┴┴┴
       `;
       expect(limitScopeByIndentation(context)(completion)).to.eq(expected);
     });
 
-    it("should allow single level closing bracket", () => {
+    it("should limit scope at next indent level, including closing line, when completion is continuing uncompleted sentence in the prefix, and starting a new indent level in next line.", () => {
+      const context = {
+        ...documentContext`
+        function findMax(arr) {║}
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+                               ├
+          let max = arr[0];
+          for (let i = 1; i < arr.length; i++) {
+            if (arr[i] > max) {
+              max = arr[i];
+            }
+          }
+          return max;
+        }
+        console.log(findMax([1, 2, 3, 4, 5]));┤
+      `;
+      const expected = inline`
+                               ├
+          let max = arr[0];
+          for (let i = 1; i < arr.length; i++) {
+            if (arr[i] > max) {
+              max = arr[i];
+            }
+          }
+          return max;
+        }┤
+      `;
+      expect(limitScopeByIndentation(context)(completion)).to.eq(expected);
+    });
+
+    it("should limit scope at next indent level, including closing line, when completion is continuing uncompleted sentence in the prefix, and starting a new indent level in next line.", () => {
+      const context = {
+        ...documentContext`
+        function findMax(arr) {
+          let max = arr[0];
+          for║
+        }
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+             ├ (let i = 1; i < arr.length; i++) {
+            if (arr[i] > max) {
+              max = arr[i];
+            }
+          }
+          return max;
+        }
+        console.log(findMax([1, 2, 3, 4, 5]));┤
+      `;
+      const expected = inline`
+             ├ (let i = 1; i < arr.length; i++) {
+            if (arr[i] > max) {
+              max = arr[i];
+            }
+          }┤
+        ┴┴
+      `;
+      expect(limitScopeByIndentation(context)(completion)).to.eq(expected);
+    });
+
+    it("should limit scope at current indent level, exclude closing line, when completion starts new sentences at same indent level.", () => {
+      const context = {
+        ...documentContext`
+        function findMax(arr) {
+          let max = arr[0];║
+        }
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+                           ├
+          for (let i = 1; i < arr.length; i++) {
+            if (arr[i] > max) {
+              max = arr[i];
+            }
+          }
+          return max;
+        }┤
+      `;
+      const expected = inline`
+                           ├
+          for (let i = 1; i < arr.length; i++) {
+            if (arr[i] > max) {
+              max = arr[i];
+            }
+          }
+          return max;┤
+        ┴┴
+      `;
+      expect(limitScopeByIndentation(context)(completion)).to.eq(expected);
+    });
+
+    it("should allow only one level closing bracket", () => {
       const context = {
         ...documentContext`
         function safeParse(json) {
@@ -58,35 +219,41 @@ describe("postprocess", () => {
       expect(limitScopeByIndentation(context)(completion)).to.eq(expected);
     });
 
-    it("should allow single level closing bracket", () => {
+    it("should allow level closing bracket at current line, it looks same as starts new sentences", () => {
       const context = {
         ...documentContext`
-        function safeParse(json) {
+        function helloworld() {
-          try {
+          console.log("hello");
-            return JSON.parse(json);
+          ║
-          } catch (e) {
+        `,
-            ║
-          }
-        }`,
         language: "javascript",
       };
       const completion = inline`
-            ├return null;
+          ├}┤
-          }
-        }┤
       `;
-      // In fact, we do not expect the closing `}`, because of there are `}` already in suffix,
+      expect(limitScopeByIndentation(context)(completion)).to.be.eq(completion);
-      // but we leave them here and pass to other filters to handle it.
-      const expected = inline`
-            ├return null;
-          }┤
-        ┴┴
-      `;
-      expect(limitScopeByIndentation(context)(completion)).to.eq(expected);
     });
 
-    // Might be better to allow
+    it("should not allow level closing bracket, when the suffix lines have same indent level", () => {
-    it("not allow back step indent level with `catch` or `else` if there is no similar block yet", () => {
+      const context = {
+        ...documentContext`
+        function helloworld() {
+          console.log("hello");║
+          console.log("world");
+        }
+        `,
+        language: "javascript",
+      };
+      const completion = inline`
+                               ├
+        }┤
+      `;
+      const expected = inline`
+                               ├┤`;
+      expect(limitScopeByIndentation(context)(completion)).to.be.eq(expected);
+    });
+
+    it("should use indent level of previous line, when current line is empty.", () => {
       const context = {
         ...documentContext`
         function safeParse(json) {
@@ -105,8 +272,11 @@ describe("postprocess", () => {
         }┤
       `;
       const expected = inline`
-            ├return JSON.parse(json);┤
+            ├return JSON.parse(json);
-        ┴┴┴┴
+          } catch (e) {
+            return null;
+          }┤
+        ┴┴
       `;
       expect(limitScopeByIndentation(context)(completion)).to.eq(expected);
     });

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

@@ -1,21 +1,8 @@
 import { PostprocessFilter, PostprocessContext, logger } from "./base";
 import { isBlank, splitLines } from "../utils";
 
-function calcIndentLevel(line) {
+function calcIndentLevel(line: string): number {
-  return line.match(/^[ \t]*/)?.[0]?.length || 0;
+  return line.match(/^[ \t]*/)?.[0]?.length ?? 0;
-}
-
-function isIndentBlockClosingAllowed(currentIndentLevel, suffixLines) {
-  let index = 1;
-  while (index < suffixLines.length && isBlank(suffixLines[index])) {
-    index++;
-  }
-  if (index >= suffixLines.length) {
-    return true;
-  } else {
-    const indentLevel = calcIndentLevel(suffixLines[index]);
-    return indentLevel < currentIndentLevel;
-  }
 }
 
 function isOpeningIndentBlock(lines, index) {
@@ -25,22 +12,112 @@ function isOpeningIndentBlock(lines, index) {
   return calcIndentLevel(lines[index]) < calcIndentLevel(lines[index + 1]);
 }
 
+function shouldOnlyAllowSingleLine(suffixLines: string[]): boolean {
+  let currentLineInSuffix = suffixLines[0] ?? "";
+  return !isBlank(currentLineInSuffix.replace(/[\)\}\]]/g, ""));
+}
+
+function processContext(
+  lines: string[],
+  prefixLines: string[],
+  suffixLines: string[],
+): { indentLevelLimit: number; allowClosingLine: boolean } {
+  let result = { indentLevelLimit: 0, allowClosingLine: false };
+  if (lines.length == 0 || prefixLines.length == 0) {
+    return result; // guard for empty input, technically unreachable
+  }
+  const currentLineInPrefix = prefixLines[prefixLines.length - 1];
+  const isCurrentLineInPrefixBlank = isBlank(currentLineInPrefix);
+  // if current line is blank, use the previous line as reference
+  let referenceLineInPrefixIndex = prefixLines.length - 1;
+  while (referenceLineInPrefixIndex >= 0 && isBlank(prefixLines[referenceLineInPrefixIndex])) {
+    referenceLineInPrefixIndex--;
+  }
+  if (referenceLineInPrefixIndex < 0) {
+    return result; // blank prefix, should be unreachable
+  }
+  const referenceLineInPrefix = prefixLines[referenceLineInPrefixIndex];
+  const referenceLineInPrefixIndent = calcIndentLevel(referenceLineInPrefix);
+
+  const currentLineInCompletion = lines[0];
+  const isCurrentLineInCompletionBlank = isBlank(currentLineInCompletion);
+  // if current line is blank, use the next line as reference
+  let referenceLineInCompletionIndex = 0;
+  while (referenceLineInCompletionIndex < lines.length && isBlank(lines[referenceLineInCompletionIndex])) {
+    referenceLineInCompletionIndex++;
+  }
+  if (referenceLineInCompletionIndex >= lines.length) {
+    return result; // blank completion, should be unreachable
+  }
+  const referenceLineInCompletion = lines[referenceLineInCompletionIndex];
+  let referenceLineInCompletionIndent;
+  if (isCurrentLineInCompletionBlank) {
+    referenceLineInCompletionIndent = calcIndentLevel(referenceLineInCompletion);
+  } else {
+    referenceLineInCompletionIndent = calcIndentLevel(currentLineInPrefix + referenceLineInCompletion);
+  }
+
+  if (!isCurrentLineInCompletionBlank && !isCurrentLineInPrefixBlank) {
+    // if two reference lines are contacted at current line, it is continuing uncompleted sentence
+
+    result.indentLevelLimit = referenceLineInPrefixIndent + 1; // + 1 for comparison, no matter how many spaces indent
+    // allow closing line if first line is opening a new indent block
+    result.allowClosingLine = !!lines[1] && calcIndentLevel(lines[1]) > referenceLineInPrefixIndent;
+  } else if (referenceLineInCompletionIndent > referenceLineInPrefixIndent) {
+    // if reference line in completion has more indent than reference line in prefix, it is opening a new indent block
+
+    result.indentLevelLimit = referenceLineInPrefixIndent + 1;
+    result.allowClosingLine = true;
+  } else if (referenceLineInCompletionIndent < referenceLineInPrefixIndent) {
+    // if reference line in completion has less indent than reference line in prefix, allow this closing
+
+    result.indentLevelLimit = referenceLineInPrefixIndent;
+    result.allowClosingLine = true;
+  } else {
+    // otherwise, it is starting a new sentence at same indent level
+
+    result.indentLevelLimit = referenceLineInPrefixIndent;
+    result.allowClosingLine = false;
+  }
+
+  // check if suffix context allows closing line
+  // skip 0 that is current line in suffix, it is processed in `shouldOnlyAllowSingleLine`
+  let firstNonBlankLineInSuffix = 1;
+  while (firstNonBlankLineInSuffix < suffixLines.length && isBlank(suffixLines[firstNonBlankLineInSuffix])) {
+    firstNonBlankLineInSuffix++;
+  }
+  if (firstNonBlankLineInSuffix < suffixLines.length) {
+    result.allowClosingLine &&= calcIndentLevel(suffixLines[firstNonBlankLineInSuffix]) < result.indentLevelLimit;
+  }
+  return result;
+}
+
 export const limitScopeByIndentation: (context: PostprocessContext) => PostprocessFilter = (context) => {
   return (input) => {
     const { prefix, suffix, prefixLines, suffixLines } = context;
     const inputLines = splitLines(input);
-    const currentIndentLevel = calcIndentLevel(prefixLines[prefixLines.length - 1]);
+    if (shouldOnlyAllowSingleLine(suffixLines)) {
+      if (inputLines.length > 1) {
+        logger.debug({ input, prefix, suffix }, "Drop content with multiple lines");
+        return null;
+      }
+    }
+    const indentContext = processContext(inputLines, prefixLines, suffixLines);
     let index;
     for (index = 1; index < inputLines.length; index++) {
       if (isBlank(inputLines[index])) {
         continue;
       }
       const indentLevel = calcIndentLevel(inputLines[index]);
-      if (indentLevel < currentIndentLevel) {
+      if (indentLevel < indentContext.indentLevelLimit) {
-        // If the line is indented less than the current indent level, it is out of scope.
+        // If the line is indented less than the indent level limit, it is closing indent block.
-        // We assume it begins with a symbol closing block.
+        // But when it is opening a new indent block immediately, such as `} else {`.
-        // If suffix context allows, and it do not open a new intent block, include this line here.
+        if (isOpeningIndentBlock(inputLines, index)) {
-        if (isIndentBlockClosingAllowed(currentIndentLevel, suffixLines) && !isOpeningIndentBlock(inputLines, index)) {
+          continue;
+        }
+        // We include this closing line here if context allows
+        // Python does not have closing bracket, so we always include closing line
+        if (indentContext.allowClosingLine && context.request.language !== "python") {
           index++;
         }
         break;