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

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 by bracket stack",
      );
    } 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 by bracket stack",
      );
    }
  }
  return response;
}