mason: fix compaction marker boundary anchoring

Author: ualtinok

packages/plugin/src/features/magic-context/compaction-marker.test.ts

@@ -0,0 +1,157 @@
+/// <reference types="bun-types" />
+
+import { afterEach, describe, expect, it } from "bun:test";
+import { mkdirSync, mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { Database } from "../../shared/sqlite";
+import { closeQuietly } from "../../shared/sqlite-helpers";
+import {
+    closeCompactionMarkerDb,
+    findBoundaryUserMessage,
+    injectCompactionMarker,
+} from "./compaction-marker";
+
+const tempDirs: string[] = [];
+const originalXdgDataHome = process.env.XDG_DATA_HOME;
+
+function useTempDataHome(prefix: string): string {
+    const dir = mkdtempSync(join(tmpdir(), prefix));
+    tempDirs.push(dir);
+    process.env.XDG_DATA_HOME = dir;
+    mkdirSync(join(dir, "opencode"), { recursive: true });
+    return dir;
+}
+
+function createOpenCodeDb(dataHome: string): Database {
+    const db = new Database(join(dataHome, "opencode", "opencode.db"));
+    db.exec("PRAGMA journal_mode=WAL");
+    db.exec(
+        "CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT, time_created INTEGER, time_updated INTEGER, data TEXT)",
+    );
+    db.exec(
+        "CREATE TABLE part (id TEXT PRIMARY KEY, message_id TEXT, session_id TEXT, time_created INTEGER, time_updated INTEGER, data TEXT)",
+    );
+    return db;
+}
+
+function insertMessage(
+    db: Database,
+    id: string,
+    role: string,
+    timeCreated: number,
+    data: Record<string, unknown> = {},
+): void {
+    db.prepare(
+        "INSERT INTO message (id, session_id, time_created, time_updated, data) VALUES (?, 'ses-1', ?, ?, ?)",
+    ).run(id, timeCreated, timeCreated, JSON.stringify({ role, ...data }));
+}
+
+afterEach(() => {
+    closeCompactionMarkerDb();
+    process.env.XDG_DATA_HOME = originalXdgDataHome;
+    for (const dir of tempDirs) {
+        rmSync(dir, { recursive: true, force: true, maxRetries: 10, retryDelay: 100 });
+    }
+    tempDirs.length = 0;
+});
+
+describe("findBoundaryUserMessage", () => {
+    it("anchors by endMessageId after rows before the target were deleted", () => {
+        const dataHome = useTempDataHome("marker-boundary-deleted-before-");
+        const db = createOpenCodeDb(dataHome);
+        insertMessage(db, "msg_001_deleted_user", "user", 100);
+        insertMessage(db, "msg_002_deleted_assistant", "assistant", 200);
+        insertMessage(db, "msg_003_prior_user", "user", 300);
+        insertMessage(db, "msg_004_target", "assistant", 400);
+        insertMessage(db, "msg_005_after_user", "user", 500);
+        db.prepare(
+            "DELETE FROM message WHERE id IN ('msg_001_deleted_user', 'msg_002_deleted_assistant')",
+        ).run();
+        closeQuietly(db);
+
+        expect(findBoundaryUserMessage("ses-1", "msg_004_target")?.id).toBe("msg_003_prior_user");
+    });
+
+    it("uses the canonical time_created/id tie-break at equal timestamps", () => {
+        const dataHome = useTempDataHome("marker-boundary-tiebreak-");
+        const db = createOpenCodeDb(dataHome);
+        insertMessage(db, "msg_a_prior_user", "user", 1_000);
+        insertMessage(db, "msg_b_target", "assistant", 1_000);
+        insertMessage(db, "msg_c_after_user", "user", 1_000);
+        closeQuietly(db);
+
+        expect(findBoundaryUserMessage("ses-1", "msg_b_target")?.id).toBe("msg_a_prior_user");
+    });
+
+    it("returns the target itself when the target message is a user", () => {
+        const dataHome = useTempDataHome("marker-boundary-target-user-");
+        const db = createOpenCodeDb(dataHome);
+        insertMessage(db, "msg_001_prior_user", "user", 100);
+        insertMessage(db, "msg_002_target_user", "user", 200);
+        closeQuietly(db);
+
+        expect(findBoundaryUserMessage("ses-1", "msg_002_target_user")?.id).toBe(
+            "msg_002_target_user",
+        );
+    });
+
+    it("is unchanged by deleting rows after the target", () => {
+        const dataHome = useTempDataHome("marker-boundary-deleted-after-");
+        const db = createOpenCodeDb(dataHome);
+        insertMessage(db, "msg_001_prior_user", "user", 100);
+        insertMessage(db, "msg_002_target", "assistant", 200);
+        insertMessage(db, "msg_003_after_user", "user", 300);
+        closeQuietly(db);
+
+        expect(findBoundaryUserMessage("ses-1", "msg_002_target")?.id).toBe("msg_001_prior_user");
+
+        const reopened = new Database(join(dataHome, "opencode", "opencode.db"));
+        reopened.prepare("DELETE FROM message WHERE id = 'msg_003_after_user'").run();
+        closeQuietly(reopened);
+        closeCompactionMarkerDb();
+
+        expect(findBoundaryUserMessage("ses-1", "msg_002_target")?.id).toBe("msg_001_prior_user");
+    });
+
+    it("finds a prior user across a long assistant/tool span", () => {
+        const dataHome = useTempDataHome("marker-boundary-long-span-");
+        const db = createOpenCodeDb(dataHome);
+        insertMessage(db, "msg_001_prior_user", "user", 100);
+        for (let i = 0; i < 150; i++) {
+            insertMessage(
+                db,
+                `msg_${String(i + 2).padStart(3, "0")}_assistant`,
+                "assistant",
+                101 + i,
+            );
+        }
+        insertMessage(db, "msg_999_target", "tool", 1_000);
+        closeQuietly(db);
+
+        expect(findBoundaryUserMessage("ses-1", "msg_999_target")?.id).toBe("msg_001_prior_user");
+    });
+});
+
+describe("injectCompactionMarker", () => {
+    it("preserves the deterministic boundary in the healthy no-deletion case", () => {
+        const dataHome = useTempDataHome("marker-inject-healthy-");
+        const db = createOpenCodeDb(dataHome);
+        insertMessage(db, "msg_001_user", "user", 100);
+        insertMessage(db, "msg_002_assistant", "assistant", 200);
+        insertMessage(db, "msg_003_target", "assistant", 300);
+        closeQuietly(db);
+
+        const result = injectCompactionMarker({
+            sessionId: "ses-1",
+            endOrdinal: 3,
+            endMessageId: "msg_003_target",
+            summaryText: "summary placeholder",
+            directory: dataHome,
+        });
+
+        // Generated marker row ids include random base62 suffixes; compare only
+        // the deterministic boundary field that the old ordinal path intended.
+        expect(result?.boundaryMessageId).toBe("msg_001_user");
+    });
+});

packages/plugin/src/features/magic-context/compaction-marker.ts

@@ -183,53 +183,99 @@ export function closeCompactionMarkerDb(): void {
 
 // ── Boundary User Message Resolution ─────────────────────────────
 
-interface BoundaryUserMessage {
+export interface BoundaryUserMessage {
     id: string;
     timeCreated: number;
 }
 
+interface NonSummaryMessageSortKey {
+    id: string;
+    timeCreated: number;
+}
+
+function getNonSummaryMessageSortKey(
+    sessionId: string,
+    messageId: string,
+): NonSummaryMessageSortKey | null {
+    const db = getWritableOpenCodeDb();
+    const row = db
+        .prepare(
+            `SELECT time_created, id
+             FROM message
+             WHERE session_id = ?
+               AND id = ?
+               AND NOT (COALESCE(json_extract(data, '$.summary'), 0) = 1
+                        AND COALESCE(json_extract(data, '$.finish'), '') = 'stop')
+             LIMIT 1`,
+        )
+        .get(sessionId, messageId) as { time_created?: unknown; id?: unknown } | undefined;
+    if (typeof row?.time_created !== "number" || typeof row.id !== "string") {
+        return null;
+    }
+    return { id: row.id, timeCreated: row.time_created };
+}
+
 /**
- * Find the nearest user message at or before the given raw ordinal.
+ * Find the nearest user message at or before the given end message id.
  * The boundary must be a user message for filterCompacted to work.
  *
  * Filters out compaction summary messages (summary=true, finish="stop")
  * so ordinals stay consistent with readRawSessionMessagesFromDb.
  */
 export function findBoundaryUserMessage(
     sessionId: string,
-    endOrdinal: number,
+    endMessageId: string,
 ): BoundaryUserMessage | null {
     const db = getWritableOpenCodeDb();
 
-    // Filter out our own injected summary messages (summary=true AND finish="stop")
-    // in SQL using json_extract — matches readRawSessionMessagesFromDb's filter so
-    // ordinals stay parity-correct. Avoids O(N) JSON.parse in JS for sessions with
-    // thousands of messages. COALESCE handles rows missing the fields.
-    const rows = db
+    // Resolve the target's canonical sort key first, using the same summary
+    // exclusion as readRawSessionMessagesFromDb. If the stored endMessageId is
+    // gone (or is itself one of our injected summaries), the pending/direct
+    // marker update is stale and must not move the boundary.
+    const target = getNonSummaryMessageSortKey(sessionId, endMessageId);
+    if (!target) return null;
+
+    // Match the raw-message reader's canonical ASC order
+    // (time_created ASC, id ASC). "At or before target" is therefore
+    // time_created < target.time_created OR the same timestamp with id <= target.id.
+    // Push role='user' into SQL so a long assistant/tool span before the target
+    // cannot exhaust a JS scan window and miss the prior user.
+    const boundary = db
         .prepare(
             `SELECT id, time_created, data
              FROM message
              WHERE session_id = ?
                AND NOT (COALESCE(json_extract(data, '$.summary'), 0) = 1
                         AND COALESCE(json_extract(data, '$.finish'), '') = 'stop')
-             ORDER BY time_created ASC, id ASC
-             LIMIT ?`,
+               AND COALESCE(json_extract(data, '$.role'), '') = 'user'
+               AND (time_created < ? OR (time_created = ? AND id <= ?))
+             ORDER BY time_created DESC, id DESC
+             LIMIT 1`,
         )
-        .all(sessionId, endOrdinal) as Array<{ id: string; time_created: number; data: string }>;
+        .get(sessionId, target.timeCreated, target.timeCreated, target.id) as
+        | { id?: unknown; time_created?: unknown; data?: unknown }
+        | undefined;
 
-    let bestMatch: BoundaryUserMessage | null = null;
-    for (const row of rows) {
-        try {
-            const info = JSON.parse(row.data);
-            if (info.role === "user") {
-                bestMatch = { id: row.id, timeCreated: row.time_created };
-            }
-        } catch {
-            // skip corrupt rows
-        }
+    if (typeof boundary?.id !== "string" || typeof boundary.time_created !== "number") {
+        return null;
     }
 
-    return bestMatch;
+    return { id: boundary.id, timeCreated: boundary.time_created };
+}
+
+export function compareOpenCodeMessagesByCanonicalOrder(
+    sessionId: string,
+    leftMessageId: string,
+    rightMessageId: string,
+): number | null {
+    const left = getNonSummaryMessageSortKey(sessionId, leftMessageId);
+    const right = getNonSummaryMessageSortKey(sessionId, rightMessageId);
+    if (!left || !right) return null;
+    if (left.timeCreated < right.timeCreated) return -1;
+    if (left.timeCreated > right.timeCreated) return 1;
+    if (left.id < right.id) return -1;
+    if (left.id > right.id) return 1;
+    return 0;
 }
 
 /**
@@ -276,10 +322,14 @@ export interface InjectCompactionMarkerArgs {
     sessionId: string;
     /** Raw ordinal of the last compartmentalized message */
     endOrdinal: number;
+    /** OpenCode message id of the last compartmentalized message */
+    endMessageId: string;
     /** Summary text for the compaction summary message (static placeholder) */
     summaryText: string;
     /** Working directory for the session */
     directory: string;
+    /** Boundary resolved before removing the old marker (prevents null-boundary cache busts). */
+    resolvedBoundary?: BoundaryUserMessage;
 }
 
 /**
@@ -298,10 +348,11 @@ export function injectCompactionMarker(
         return null;
     }
 
-    const boundary = findBoundaryUserMessage(args.sessionId, args.endOrdinal);
+    const boundary =
+        args.resolvedBoundary ?? findBoundaryUserMessage(args.sessionId, args.endMessageId);
     if (!boundary) {
         log(
-            `[magic-context] compaction-marker: no user message found at or before ordinal ${args.endOrdinal}`,
+            `[magic-context] compaction-marker: no user message found at or before endMessageId ${args.endMessageId} (ordinal ${args.endOrdinal})`,
         );
         return null;
     }

packages/plugin/src/features/magic-context/migrations-v39.test.ts

@@ -0,0 +1,75 @@
+/// <reference types="bun-types" />
+
+import { describe, expect, test } from "bun:test";
+import { Database } from "../../shared/sqlite";
+import { closeQuietly } from "../../shared/sqlite-helpers";
+import { runMigrations } from "./migrations";
+import { initializeDatabase } from "./storage-db";
+
+function columnNames(db: Database): string[] {
+    return (db.prepare("PRAGMA table_info(session_meta)").all() as Array<{ name: string }>).map(
+        (column) => column.name,
+    );
+}
+
+describe("migration v39 — compaction marker target end message id", () => {
+    test("fresh database has nullable compaction_marker_target_end_message_id column", () => {
+        const db = new Database(":memory:");
+        try {
+            initializeDatabase(db);
+            runMigrations(db);
+
+            expect(columnNames(db)).toContain("compaction_marker_target_end_message_id");
+            db.prepare("INSERT INTO session_meta (session_id) VALUES ('ses-1')").run();
+            const row = db
+                .prepare(
+                    "SELECT compaction_marker_target_end_message_id FROM session_meta WHERE session_id = 'ses-1'",
+                )
+                .get() as { compaction_marker_target_end_message_id: string | null };
+            expect(row.compaction_marker_target_end_message_id).toBeNull();
+        } finally {
+            closeQuietly(db);
+        }
+    });
+
+    test("upgrade adds column and backfills JSON targetEndMessageId when present", () => {
+        const db = new Database(":memory:");
+        try {
+            db.exec(`
+                CREATE TABLE schema_migrations (version INTEGER PRIMARY KEY, description TEXT NOT NULL, applied_at INTEGER NOT NULL);
+                INSERT INTO schema_migrations (version, description, applied_at) VALUES (38, 'pre-v39 fixture', 1);
+                CREATE TABLE session_meta (
+                    session_id TEXT PRIMARY KEY,
+                    compaction_marker_state TEXT DEFAULT ''
+                );
+                INSERT INTO session_meta (session_id, compaction_marker_state) VALUES (
+                    'ses-json',
+                    '{"boundaryMessageId":"msg-user","summaryMessageId":"msg-summary","compactionPartId":"prt-comp","summaryPartId":"prt-summary","boundaryOrdinal":10,"targetEndMessageId":"msg-target"}'
+                );
+                INSERT INTO session_meta (session_id, compaction_marker_state) VALUES (
+                    'ses-legacy',
+                    '{"boundaryMessageId":"msg-user","summaryMessageId":"msg-summary","compactionPartId":"prt-comp","summaryPartId":"prt-summary","boundaryOrdinal":10}'
+                );
+            `);
+
+            runMigrations(db);
+            runMigrations(db);
+
+            expect(columnNames(db)).toContain("compaction_marker_target_end_message_id");
+            const rows = db
+                .prepare(
+                    "SELECT session_id, compaction_marker_target_end_message_id FROM session_meta ORDER BY session_id",
+                )
+                .all() as Array<{
+                session_id: string;
+                compaction_marker_target_end_message_id: string | null;
+            }>;
+            expect(rows).toEqual([
+                { session_id: "ses-json", compaction_marker_target_end_message_id: "msg-target" },
+                { session_id: "ses-legacy", compaction_marker_target_end_message_id: null },
+            ]);
+        } finally {
+            closeQuietly(db);
+        }
+    });
+});