refactor(agent): extract calculateReplaceRange and add unit test. (#764)

* refactor(agent): extract calculateReplaceRange.

* test(agent): add unit test for calculateReplaceRangeByBracketStack.

clients/tabby-agent/src/TabbyAgent.ts

@@ -5,14 +5,7 @@ import { deepmerge } from "deepmerge-ts";
 import { getProperty, setProperty, deleteProperty } from "dot-prop";
 import createClient from "openapi-fetch";
 import { paths as TabbyApi } from "./types/tabbyApi";
-import {
-  isBlank,
-  abortSignalFromAnyOf,
-  findUnpairedAutoClosingChars,
-  HttpError,
-  isTimeoutError,
-  isCanceledError,
-} from "./utils";
+import { isBlank, abortSignalFromAnyOf, HttpError, isTimeoutError, isCanceledError } from "./utils";
 import type {
   Agent,
   AgentStatus,
@@ -32,7 +25,7 @@ import { CompletionCache } from "./CompletionCache";
 import { CompletionDebounce } from "./CompletionDebounce";
 import { CompletionContext } from "./CompletionContext";
 import { DataStore } from "./dataStore";
-import { preCacheProcess, postCacheProcess } from "./postprocess";
+import { preCacheProcess, postCacheProcess, calculateReplaceRange } from "./postprocess";
 import { rootLogger, allLoggers } from "./logger";
 import { AnonymousUsageLogger } from "./AnonymousUsageLogger";
 import { CompletionProviderStats, CompletionProviderStatsEntry } from "./CompletionProviderStats";
@@ -291,35 +284,6 @@ export class TabbyAgent extends EventEmitter implements Agent {
     return { prefix, suffix };
   }
 
-  private calculateReplaceRange(response: CompletionResponse, context: CompletionContext): CompletionResponse {
-    const { suffixLines } = context;
-    const suffixText = suffixLines[0]?.trimEnd() || "";
-    if (isBlank(suffixText)) {
-      return response;
-    }
-    for (const choice of response.choices) {
-      const completionText = choice.text.slice(context.position - choice.replaceRange.start);
-      const unpaired = findUnpairedAutoClosingChars(completionText);
-      if (isBlank(unpaired)) {
-        continue;
-      }
-      if (suffixText.startsWith(unpaired)) {
-        choice.replaceRange.end = context.position + unpaired.length;
-        this.logger.trace(
-          { context, completion: choice.text, range: choice.replaceRange, unpaired },
-          "Adjust replace range",
-        );
-      } else if (unpaired.startsWith(suffixText)) {
-        choice.replaceRange.end = context.position + suffixText.length;
-        this.logger.trace(
-          { context, completion: choice.text, range: choice.replaceRange, unpaired },
-          "Adjust replace range",
-        );
-      }
-    }
-    return response;
-  }
-
   public async initialize(options: AgentInitOptions): Promise<boolean> {
     if (options.clientProperties) {
       const { user: userProp, session: sessionProp } = options.clientProperties;
@@ -574,7 +538,10 @@ export class TabbyAgent extends EventEmitter implements Agent {
         throw options.signal.reason;
       }
       // Calculate replace range
-      completionResponse = this.calculateReplaceRange(completionResponse, context);
+      completionResponse = await calculateReplaceRange(completionResponse, context);
+      if (options?.signal?.aborted) {
+        throw options.signal.reason;
+      }
     } catch (error) {
       if (isCanceledError(error) || isTimeoutError(error)) {
         if (stats) {

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

@@ -0,0 +1,213 @@
+import { expect } from "chai";
+import { documentContext, inline } from "./testUtils";
+import { calculateReplaceRangeByBracketStack } from "./calculateReplaceRangeByBracketStack";
+
+describe("postprocess", () => {
+  describe("calculateReplaceRangeByBracketStack", () => {
+    it("should handle auto closing quotes", () => {
+      const context = {
+        ...documentContext`
+        const hello = "║"
+        `,
+        language: "typescript",
+      };
+      const response = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                       ├hello";┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position,
+            },
+          },
+        ],
+      };
+      const expected = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                       ├hello";┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position + 1,
+            },
+          },
+        ],
+      };
+      expect(calculateReplaceRangeByBracketStack(response, context)).to.deep.equal(expected);
+    });
+
+    it("should handle auto closing quotes", () => {
+      const context = {
+        ...documentContext`
+        let htmlMarkup = \`║\`
+        `,
+        language: "typescript",
+      };
+      const response = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                           ├<h1>\${message}</h1>\`;┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position,
+            },
+          },
+        ],
+      };
+      const expected = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                           ├<h1>\${message}</h1>\`;┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position + 1,
+            },
+          },
+        ],
+      };
+      expect(calculateReplaceRangeByBracketStack(response, context)).to.deep.equal(expected);
+    });
+
+    it("should handle multiple auto closing brackets", () => {
+      const context = {
+        ...documentContext`
+        process.on('data', (data) => {║})
+        `,
+        language: "typescript",
+      };
+      const response = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                                      ├
+              console.log(data);
+            });┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position,
+            },
+          },
+        ],
+      };
+      const expected = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                                      ├
+              console.log(data);
+            });┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position + 2,
+            },
+          },
+        ],
+      };
+      expect(calculateReplaceRangeByBracketStack(response, context)).to.deep.equal(expected);
+    });
+
+    it("should handle multiple auto closing brackets", () => {
+      const context = {
+        ...documentContext`
+        let mat: number[][][] = [[[║]]]
+        `,
+        language: "typescript",
+      };
+      const response = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                                   ├1, 2], [3, 4]], [[5, 6], [7, 8]]];┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position,
+            },
+          },
+        ],
+      };
+      const expected = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                                   ├1, 2], [3, 4]], [[5, 6], [7, 8]]];┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position + 3,
+            },
+          },
+        ],
+      };
+      expect(calculateReplaceRangeByBracketStack(response, context)).to.deep.equal(expected);
+    });
+  });
+
+  describe("calculateReplaceRangeByBracketStack: bad cases", () => {
+    const context = {
+      ...documentContext`
+      function clamp(n: number, max: number, min: number): number {
+        return Math.max(Math.min(║);
+      }
+      `,
+      language: "typescript",
+    };
+    const response = {
+      id: "",
+      choices: [
+        {
+          index: 0,
+          text: inline`
+                                 ├n, max), min┤
+          `,
+          replaceRange: {
+            start: context.position,
+            end: context.position,
+          },
+        },
+      ],
+    };
+    const expected = {
+      id: "",
+      choices: [
+        {
+          index: 0,
+          text: inline`
+                                 ├n, max), min┤
+          `,
+          replaceRange: {
+            start: context.position,
+            end: context.position,
+          },
+        },
+      ],
+    };
+    expect(calculateReplaceRangeByBracketStack(response, context)).not.to.deep.equal(expected);
+  });
+});

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

@@ -0,0 +1,29 @@
+import { CompletionContext, CompletionResponse } from "../Agent";
+import { isBlank, findUnpairedAutoClosingChars } from "../utils";
+import { logger } from "./base";
+
+export function calculateReplaceRangeByBracketStack(
+  response: CompletionResponse,
+  context: CompletionContext,
+): CompletionResponse {
+  const { suffixLines } = context;
+  const suffixText = suffixLines[0]?.trimEnd() || "";
+  if (isBlank(suffixText)) {
+    return response;
+  }
+  for (const choice of response.choices) {
+    const completionText = choice.text.slice(context.position - choice.replaceRange.start);
+    const unpaired = findUnpairedAutoClosingChars(completionText);
+    if (isBlank(unpaired)) {
+      continue;
+    }
+    if (suffixText.startsWith(unpaired)) {
+      choice.replaceRange.end = context.position + unpaired.length;
+      logger.trace({ context, completion: choice.text, range: choice.replaceRange, unpaired }, "Adjust replace range");
+    } else if (unpaired.startsWith(suffixText)) {
+      choice.replaceRange.end = context.position + suffixText.length;
+      logger.trace({ context, completion: choice.text, range: choice.replaceRange, unpaired }, "Adjust replace range");
+    }
+  }
+  return response;
+}

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

@@ -8,6 +8,7 @@ import { limitScopeByIndentation } from "./limitScopeByIndentation";
 import { trimSpace } from "./trimSpace";
 import { dropDuplicated } from "./dropDuplicated";
 import { dropBlank } from "./dropBlank";
+import { calculateReplaceRangeByBracketStack } from "./calculateReplaceRangeByBracketStack";
 
 export async function preCacheProcess(
   context: CompletionContext,
@@ -34,3 +35,10 @@ export async function postCacheProcess(
     .then(applyFilter(trimSpace(context), context))
     .then(applyFilter(dropBlank(), context));
 }
+
+export async function calculateReplaceRange(
+  response: CompletionResponse,
+  context: CompletionContext,
+): Promise<CompletionResponse> {
+  return calculateReplaceRangeByBracketStack(response, context);
+}

clients/tabby-agent/tests/bad_cases/2-typescript.ts

@@ -0,0 +1,3 @@
+function clamp(n: number, max: number, min: number): number {
+  return Math.max(Math.min(⏩⏭n, max), min⏮⏪);
+}