From 50973613471140b52353f8e1e1c21c39c8501d9c Mon Sep 17 00:00:00 2001
From: Dex Hunter <i@dex.moe>
Date: Mon, 13 Jul 2026 04:29:05 +0000
Subject: [PATCH] perf(llm): avoid tokenizing discarded lines

Fixes #12980.
---
 core/llm/countTokens.test.ts |  4 +--
 core/llm/countTokens.ts      | 55 +++++++++++++++---------------------
 2 files changed, 25 insertions(+), 34 deletions(-)

diff --git a/core/llm/countTokens.test.ts b/core/llm/countTokens.test.ts
index fa4592cdc66..0327dffcf2f 100644
--- a/core/llm/countTokens.test.ts
+++ b/core/llm/countTokens.test.ts
@@ -47,7 +47,7 @@ describe("countTokensAsync", () => {
   });
 });
 
-describe.skip("pruneLinesFromTop", () => {
+describe("pruneLinesFromTop", () => {
   it("should prune lines from the top to fit within max tokens", () => {
     const prompt = "Line 1\nLine 2\nLine 3\nLine 4";
     const pruned = pruneLinesFromTop(prompt, 5, "gpt-4");
@@ -88,7 +88,7 @@ describe.skip("pruneLinesFromTop", () => {
   });
 });
 
-describe.skip("pruneLinesFromBottom", () => {
+describe("pruneLinesFromBottom", () => {
   it("should prune lines from the bottom to fit within max tokens", () => {
     const prompt = "Line 1\nLine 2\nLine 3\nLine 4";
     const pruned = pruneLinesFromBottom(prompt, 5, "gpt-4");
diff --git a/core/llm/countTokens.ts b/core/llm/countTokens.ts
index b742d70b0f0..03166e4a30b 100644
--- a/core/llm/countTokens.ts
+++ b/core/llm/countTokens.ts
@@ -285,24 +285,19 @@ function pruneLinesFromTop(
   modelName: string,
 ): string {
   const lines = prompt.split("\n");
-  // Preprocess tokens for all lines and cache them.
-  const lineTokens = lines.map((line) => countTokens(line, modelName));
-  let totalTokens = lineTokens.reduce((sum, tokens) => sum + tokens, 0);
-  let start = 0;
-  let currentLines = lines.length;
-
-  // Calculate initial token count including newlines
-  totalTokens += Math.max(0, currentLines - 1); // Add tokens for joining newlines
-
-  // Using indexes instead of array modifications.
-  // Remove lines from the top until the token count is within the limit.
-  while (totalTokens > maxTokens && start < currentLines) {
-    totalTokens -= lineTokens[start];
-    // Decrement token count for the removed line and its preceding/joining newline (if not the last line)
-    if (currentLines - start > 1) {
-      totalTokens--;
+  let start = lines.length;
+  let totalTokens = 0;
+
+  // Build the retained suffix without tokenizing discarded leading lines.
+  while (start > 0) {
+    const next = start - 1;
+    const nextTokens = countTokens(lines[next], modelName);
+    const newlineTokens = start < lines.length ? 1 : 0;
+    if (totalTokens + nextTokens + newlineTokens > maxTokens) {
+      break;
     }
-    start++;
+    totalTokens += nextTokens + newlineTokens;
+    start = next;
   }
 
   return lines.slice(start).join("\n");
@@ -314,22 +309,18 @@ function pruneLinesFromBottom(
   modelName: string,
 ): string {
   const lines = prompt.split("\n");
-  const lineTokens = lines.map((line) => countTokens(line, modelName));
-  let totalTokens = lineTokens.reduce((sum, tokens) => sum + tokens, 0);
-  let end = lines.length;
-
-  // Calculate initial token count including newlines
-  totalTokens += Math.max(0, end - 1); // Add tokens for joining newlines
-
-  // Reverse traversal to avoid array modification
-  // Remove lines from the bottom until the token count is within the limit.
-  while (totalTokens > maxTokens && end > 0) {
-    end--;
-    totalTokens -= lineTokens[end];
-    // Decrement token count for the removed line and its following/joining newline (if not the first line)
-    if (end > 0) {
-      totalTokens--;
+  let end = 0;
+  let totalTokens = 0;
+
+  // Build the retained prefix without tokenizing discarded trailing lines.
+  while (end < lines.length) {
+    const nextTokens = countTokens(lines[end], modelName);
+    const newlineTokens = end > 0 ? 1 : 0;
+    if (totalTokens + nextTokens + newlineTokens > maxTokens) {
+      break;
     }
+    totalTokens += nextTokens + newlineTokens;
+    end++;
   }
 
   return lines.slice(0, end).join("\n");