Merge commit '667dbe0b6b3e'

Co-authored-by: Alfonso <alfonso-magic-context@users.noreply.github.com>

Author: ualtinok

packages/plugin/src/features/magic-context/protection-window-paged.test.ts

@@ -0,0 +1,115 @@
+import { describe, expect, it } from "bun:test";
+import { Database } from "../../shared/sqlite";
+import {
+    computeProtectionWindow,
+    getProtectionWindowForSession,
+    type ProtectionWindowRow,
+} from "./protection-window";
+import { TAG_SELECT_COLUMNS } from "./storage-tags";
+
+type SeedRow = { tag: number; tokens: number | null; type?: string };
+const generated = (count: number): SeedRow[] =>
+    Array.from({ length: count }, (_, i) => ({
+        tag: i + 1,
+        tokens: i % 7 === 0 ? null : i % 5 === 0 ? 0 : 100,
+    }));
+
+function compareWithReference(seeds: SeedRow[], floor: number) {
+    const db = new Database(":memory:");
+    try {
+        // Legacy chronology can contain duplicate tag numbers; the current migration
+        // adds uniqueness, so use the reader's older supported row shape for tie coverage.
+        db.exec(`CREATE TABLE tags (
+            id INTEGER PRIMARY KEY, session_id TEXT, message_id TEXT, type TEXT,
+            status TEXT, drop_mode TEXT, tool_name TEXT, input_byte_size INTEGER DEFAULT 0,
+            byte_size INTEGER, reasoning_byte_size INTEGER DEFAULT 0,
+            tag_number INTEGER NOT NULL, caveman_depth INTEGER DEFAULT 0,
+            tool_owner_message_id TEXT, token_count INTEGER
+        ); CREATE INDEX tags_chronology ON tags(session_id, tag_number, id)`);
+        const insert = db.prepare(`INSERT INTO tags
+            (session_id, message_id, type, status, tag_number, token_count, byte_size,
+             tool_name, tool_owner_message_id, drop_mode)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`);
+        db.transaction(() => {
+            seeds.forEach((row, i) => {
+                insert.run(
+                    "paged",
+                    `message-${i}`,
+                    row.type ?? "tool",
+                    i % 2 === 0 ? "dropped" : "active",
+                    row.tag,
+                    row.tokens,
+                    i * 13,
+                    "read",
+                    `owner-${i}`,
+                    i % 2 === 0 ? "skeleton" : null,
+                );
+            });
+            insert.run(
+                "other-session",
+                "other",
+                "tool",
+                "active",
+                999999,
+                99999,
+                1,
+                "read",
+                "other",
+                null,
+            );
+        })();
+        const allRows = db
+            .prepare(`SELECT ${TAG_SELECT_COLUMNS} FROM tags
+            WHERE session_id = ? AND type = 'tool' ORDER BY tag_number ASC, id ASC`)
+            .all("paged") as ProtectionWindowRow[];
+        const expected = computeProtectionWindow(allRows, floor);
+        const actual = getProtectionWindowForSession(db, "paged", floor);
+        const { isProtected: expectedPredicate, ...expectedFields } = expected;
+        const { isProtected: actualPredicate, ...actualFields } = actual;
+        expect(actualFields).toEqual(expectedFields);
+        for (const row of allRows) expect(actualPredicate(row)).toBe(expectedPredicate(row));
+        return actual;
+    } finally {
+        db.close();
+    }
+}
+
+describe("paged protection-window differential", () => {
+    it("completes a duplicate tag group across the 256-row page boundary", () => {
+        const seeds = generated(27);
+        for (let i = 0; i < 300; i++) seeds.push({ tag: 28, tokens: 100 });
+        seeds.push({ tag: 29, tokens: 100 }, { tag: 30, tokens: 100 });
+        const result = compareWithReference(seeds, 100);
+        expect(result.status.protectedCount).toBe(302);
+        expect(result.cutoff).toBe(28);
+    });
+
+    for (const count of [0, 1, 2, 256, 257, 20000]) {
+        for (const floor of [0, 4000, count * 100 + 1]) {
+            it(`matches all fields for ${count} rows at floor ${floor}`, () => {
+                compareWithReference(generated(count), floor);
+            });
+        }
+    }
+
+    it("exhausts zero and NULL mass history with interleaved dropped rows", () => {
+        compareWithReference(
+            generated(600).map((row, i) => ({ ...row, tokens: i % 2 ? 0 : null })),
+            4000,
+        );
+    });
+
+    it("retains complete ties with fewer than three distinct tool tags", () => {
+        compareWithReference(
+            generated(600).map((row, i) => ({ ...row, tag: i % 2 })),
+            0,
+        );
+    });
+
+    it("ignores non-tool chronology and other sessions", () => {
+        compareWithReference(
+            generated(600).map((row, i) => ({ ...row, type: i % 3 ? "tool" : "message" })),
+            4000,
+        );
+    });
+});

packages/plugin/src/features/magic-context/protection-window.ts

@@ -261,8 +261,9 @@ export function readEpochFloorSnapshot(db: Database, sessionId: string): number
 }
 
 /**
- * Compute the protection window for a session by loading all persisted tool rows
- * from the database.
+ * Read only the persisted tool suffix needed for protection. A tag-number group
+ * can span pages, so test the stopping condition only when the next group starts.
+ * Dropped rows still contribute to chronology and mass, just as in the pure walk.
  */
 export function getProtectionWindowForSession(
     db: Database,
@@ -274,13 +275,37 @@ export function getProtectionWindowForSession(
             ? floor
             : (readEpochFloorSnapshot(db, sessionId) ?? 0);
 
-    const rows = db
-        .prepare(
-            `SELECT ${TAG_SELECT_COLUMNS} FROM tags
-             WHERE session_id = ? AND type = 'tool'
-             ORDER BY tag_number ASC, id ASC`,
-        )
-        .all(sessionId) as ProtectionWindowRow[];
+    const pageSize = 256;
+    // Exclude the tool-owner partial index from planning: without ANALYZE it can
+    // win over chronology and force a whole-session sort before each small page.
+    const projection = `SELECT ${TAG_SELECT_COLUMNS} FROM tags
+        WHERE session_id = ? AND +type = 'tool'`;
+    const order = `ORDER BY tag_number DESC, id DESC LIMIT ${pageSize}`;
+    let page = db.prepare(`${projection} ${order}`).all(sessionId) as ProtectionWindowRow[];
+    const rows: ProtectionWindowRow[] = [];
+    let mass = 0;
+    let distinctTags = 0;
+    let currentTag: number | null = null;
+
+    while (page.length > 0) {
+        for (const row of page) {
+            const tagNumber = getRowTagNumber(row);
+            if (tagNumber !== currentTag) {
+                if (distinctTags >= 3 && mass >= effectiveFloor) {
+                    return computeProtectionWindow(rows, effectiveFloor);
+                }
+                distinctTags++;
+                currentTag = tagNumber;
+            }
+            rows.push(row);
+            mass += rowWindowMass(row);
+        }
+        if (page.length < pageSize) break;
+        const last = page[page.length - 1];
+        page = db
+            .prepare(`${projection} AND (tag_number, id) < (?, ?) ${order}`)
+            .all(sessionId, getRowTagNumber(last), getRowIdentity(last)) as ProtectionWindowRow[];
+    }
 
     return computeProtectionWindow(rows, effectiveFloor);
 }