From 01e8dea6c5d9a2b69ff12e55f1189e1e49711d67 Mon Sep 17 00:00:00 2001
From: Victor <70475442+vsolano9@users.noreply.github.com>
Date: Sun, 23 Aug 2026 16:26:33 +0200
Subject: [PATCH] perf: make budget trimming linear

Measure the transformed payload once and subtract each removed message token count from a running total instead of recounting the shrinking payload on every loop iteration. Add a 500-message regression that pins output, report arithmetic, and linear counter work.

Closes AgentPostmortem/tokencut#3
---
 src/index.mjs       | 10 ++++++++--
 test/basic.test.mjs | 24 ++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/index.mjs b/src/index.mjs
index b1c1fd4..f8cc850 100644
--- a/src/index.mjs
+++ b/src/index.mjs
@@ -107,23 +107,29 @@ export function compact(payload, opts = {}) {
     m.content = kept;
   }
 
+  // Measure the transformed payload once, then keep the total additive while
+  // trimming so each dropped message is counted only once.
+  let after = analyzePayload(out, { counter }).totalTokens;
+
   // trim oldest non-system messages to fit a hard budget
   if (maxTokens != null) {
     const isSystem = (m) => (m.role || "") === "system";
     let i = 0;
-    while (analyzePayload(out, { counter }).totalTokens > maxTokens) {
+    while (after > maxTokens) {
       const list = Array.isArray(out) ? out : out.messages;
       const trimmableEnd = list.length - keepLastTurns;
       // find the oldest non-system, trimmable message
       let idx = -1;
       for (let j = 0; j < Math.max(0, trimmableEnd); j++) { if (!isSystem(list[j])) { idx = j; break; } }
       if (idx < 0) break; // nothing left safe to drop
+      const removed = list[idx];
+      after -= blocksOf(removed.content, removed.role || "user")
+        .reduce((total, unit) => total + counter(unit.text), 0);
       list.splice(idx, 1);
       actions.push("drop:oldest-message");
       if (++i > 1000) break;
     }
   }
 
-  const after = analyzePayload(out, { counter }).totalTokens;
   return { payload: out, report: { beforeTokens: before, afterTokens: after, savedTokens: before - after, savedPct: before ? Math.round(((before - after) / before) * 100) : 0, actions } };
 }
diff --git a/test/basic.test.mjs b/test/basic.test.mjs
index f76e872..98d1bdc 100644
--- a/test/basic.test.mjs
+++ b/test/basic.test.mjs
@@ -50,3 +50,27 @@ test("compact trims oldest messages to a hard budget, keeps system", () => {
   assert.ok(report.afterTokens <= 300 || report.savedTokens > 0);
   assert.equal(out.messages[0].role, "system");
 });
+
+test("compact trims a large payload with linear token counting", () => {
+  const messages = Array.from({ length: 500 }, (_, i) => ({
+    role: "user",
+    content: `message-${i}`,
+  }));
+  let counterCalls = 0;
+  const counter = () => {
+    counterCalls++;
+    return 1;
+  };
+
+  const { payload: out, report } = compact(
+    { messages },
+    { maxTokens: 4, keepLastTurns: 4, dropDuplicates: false, counter },
+  );
+
+  assert.deepEqual(out.messages, messages.slice(-4));
+  assert.equal(report.beforeTokens, 500);
+  assert.equal(report.afterTokens, 4);
+  assert.equal(report.savedTokens, 496);
+  assert.equal(report.actions.length, 496);
+  assert.ok(counterCalls <= messages.length * 3, `counter called ${counterCalls} times`);
+});