mason: fix historian wrapup coordination

Co-authored-by: Alfonso [Magic Context] <288211368+alfonso-magic-context@users.noreply.github.com>

Author: ualtinok

packages/pi-plugin/src/commands/ctx-wrapup.test.ts

@@ -1,6 +1,10 @@
 /// <reference types="bun-types" />
 
 import { describe, expect, it, mock } from "bun:test";
+import {
+	acquireCompartmentLease,
+	releaseCompartmentLease,
+} from "@magic-context/core/features/magic-context/compartment-lease";
 import {
 	appendCompartments,
 	getCompartments,
@@ -241,6 +245,37 @@ describe("Pi /ctx-wrapup", () => {
 		}
 	});
 
+	it("bounds waiting for a foreign compartment lease and clears the wrapup marker", async () => {
+		const db = createDb();
+		try {
+			const sessionId = "pi-wrapup-lease-timeout";
+			const foreignHolder = "foreign-lease-holder";
+			expect(
+				acquireCompartmentLease(db, sessionId, foreignHolder),
+			).not.toBeNull();
+			const runPiHistorianForWrapup = mock(async () => {});
+
+			const result = await runPiWrapup(
+				pi().api,
+				deps(db, {
+					runPiHistorianForWrapup,
+					wrapupLeaseWaitTimeoutMs: 0,
+				}),
+				ctx(sessionId, 8),
+				sessionId,
+				2,
+			);
+
+			expect(result).toContain("## Magic Wrapup — Partial");
+			expect(result).toContain("Timed out waiting");
+			expect(runPiHistorianForWrapup).not.toHaveBeenCalled();
+			expect(getWrapupInProgressState(db, sessionId)).toBeNull();
+			releaseCompartmentLease(db, sessionId, foreignHolder);
+		} finally {
+			closeQuietly(db);
+		}
+	});
+
 	it("signals deferred history and materialization after a wrapup publish", async () => {
 		const db = createDb();
 		try {

packages/pi-plugin/src/commands/ctx-wrapup.ts

@@ -60,6 +60,7 @@ export interface RegisterCtxWrapupDeps {
 		[modelKey: string]: number | undefined;
 	};
 	runPiHistorianForWrapup?: typeof runPiHistorian;
+	wrapupLeaseWaitTimeoutMs?: number;
 	resolveRuntimeDeps?: (ctx: { cwd: string }) => CtxWrapupRuntimeDeps;
 }
 
@@ -70,6 +71,18 @@ export type CtxWrapupRuntimeDeps = Omit<
 
 const DEFAULT_MESSAGES_TO_KEEP = 20;
 const LEASE_WAIT_MS = 1_000;
+const MAX_WRAPUP_LEASE_WAIT_MS = 10 * 60 * 1_000;
+
+type LeaseAcquireResult =
+	| { ok: true; holderId: string }
+	| { ok: false; reason: "ownership_lost" | "timeout" };
+
+function resolveWrapupLeaseWaitTimeout(deps: CtxWrapupRuntimeDeps): number {
+	const configured = deps.wrapupLeaseWaitTimeoutMs ?? MAX_WRAPUP_LEASE_WAIT_MS;
+	return Number.isFinite(configured) && configured >= 0
+		? configured
+		: MAX_WRAPUP_LEASE_WAIT_MS;
+}
 
 export function parseWrapupArgs(
 	raw: string,
@@ -345,17 +358,21 @@ export async function runPiWrapup(
 					level: "info",
 				});
 
-				const leaseHolder = await acquireCompartmentLeaseEventually(
+				const leaseResult = await acquireCompartmentLeaseEventually(
 					deps.db,
 					sessionId,
 					renewWrapupMarker,
+					resolveWrapupLeaseWaitTimeout(deps),
 				);
-				if (!leaseHolder) {
+				if (!leaseResult.ok) {
 					failure = ownershipLost
 						? `${ownershipLostReason}; wrapped up through message ${lastEnd}. Run /ctx-wrapup again to continue.`
-						: "Another Magic Context rebuild started while wrapup was waiting. Run /ctx-wrapup again to continue.";
+						: leaseResult.reason === "timeout"
+							? `Timed out waiting for another process to release the compartment-state lease; wrapped up through message ${lastEnd}. Run /ctx-wrapup again to continue.`
+							: "Another Magic Context rebuild started while wrapup was waiting. Run /ctx-wrapup again to continue.";
 					break;
 				}
+				const leaseHolder = leaseResult.holderId;
 				const leaseRenewal = setInterval(() => {
 					try {
 						renewCompartmentLease(deps.db, sessionId, leaseHolder);
@@ -480,13 +497,20 @@ async function acquireCompartmentLeaseEventually(
 	renewWrapupMarker: (
 		updates: Parameters<typeof updateWrapupInProgress>[3],
 	) => boolean,
-): Promise<string | null> {
+	maxWaitMs: number,
+): Promise<LeaseAcquireResult> {
+	const waitStartedAt = Date.now();
+	const remainingMs = (): number =>
+		Math.max(0, waitStartedAt + maxWaitMs - Date.now());
 	for (;;) {
+		if (remainingMs() <= 0) return { ok: false, reason: "timeout" };
 		const holderId = crypto.randomUUID();
 		const lease = acquireCompartmentLease(db, sessionId, holderId);
-		if (lease) return holderId;
-		if (!renewWrapupMarker({})) return null;
-		await new Promise((resolve) => setTimeout(resolve, LEASE_WAIT_MS));
+		if (lease) return { ok: true, holderId };
+		if (!renewWrapupMarker({})) return { ok: false, reason: "ownership_lost" };
+		await new Promise((resolve) =>
+			setTimeout(resolve, Math.min(LEASE_WAIT_MS, remainingMs())),
+		);
 	}
 }
 

packages/plugin/src/hooks/magic-context/compartment-runner.test.ts

@@ -34,6 +34,7 @@ import { Database } from "../../shared/sqlite";
 import { closeQuietly } from "../../shared/sqlite-helpers";
 import {
     executeContextRecomp,
+    executeContextRecompWithResult,
     getActiveCompartmentRun,
     registerActiveCompartmentRun,
     runCompartmentAgent,
@@ -175,6 +176,68 @@ describe("executeContextRecomp", () => {
         expect(getSessionFacts(db, "ses-recomp")).toHaveLength(0);
     });
 
+    it("aborts after winning the recomp lease if a wrapup marker appeared meanwhile", async () => {
+        useTempDataHome("magic-recomp-wrapup-race-");
+        const sessionId = "ses-recomp-wrapup-race";
+        createOpenCodeDb(sessionId, [
+            { id: "m-1", role: "user", text: "eligible one" },
+            { id: "m-2", role: "assistant", text: "eligible two" },
+            { id: "m-3", role: "user", text: "protected 1" },
+            { id: "m-4", role: "user", text: "protected 2" },
+            { id: "m-5", role: "user", text: "protected 3" },
+            { id: "m-6", role: "user", text: "protected 4" },
+            { id: "m-7", role: "user", text: "protected 5" },
+        ]);
+        const db = openDatabase();
+        getOrCreateSessionMeta(db, sessionId);
+
+        const client = {
+            session: {
+                get: mock(async () => {
+                    throw new Error("recomp runner should not start under a live wrapup");
+                }),
+                create: mock(async () => ({ data: { id: "unused" } })),
+                prompt: mock(async () => ({})),
+                messages: mock(async () => ({ data: [] })),
+                delete: mock(async () => ({})),
+            },
+        } as unknown as PluginContext["client"];
+
+        const result = await executeContextRecompWithResult(
+            {
+                client,
+                db,
+                sessionId,
+                historianChunkTokens: 10_000,
+                directory: "/tmp",
+            },
+            {
+                onLeaseAcquired: () => {
+                    const acquired = acquireWrapupInProgress(db, sessionId, {
+                        holderId: "wrapup-holder",
+                        messagesToKeep: 5,
+                        anchorRawMessageCount: 7,
+                        targetEligibleEndOrdinal: 3,
+                        lastCompartmentEnd: -1,
+                        chunkIndex: 0,
+                        expectedChunks: 1,
+                    });
+                    expect(acquired.ok).toBe(true);
+                },
+            },
+        );
+
+        expect(result.published).toBe(false);
+        expect(result.message).toContain("/ctx-wrapup is already compacting");
+        expect(client.session.get).not.toHaveBeenCalled();
+        expect(getCompartments(db, sessionId)).toHaveLength(0);
+        expect(
+            db
+                .prepare("SELECT 1 FROM compartment_state_lease WHERE session_id = ?")
+                .get(sessionId) ?? null,
+        ).toBeNull();
+    });
+
     it("keeps published state unchanged when a later recomp pass fails", async () => {
         useTempDataHome("magic-recomp-fail-closed-");
         createOpenCodeDb("ses-recomp-fail", [

packages/plugin/src/hooks/magic-context/compartment-runner.ts

@@ -206,6 +206,8 @@ export interface ExecuteContextRecompOptions {
      * replacing all compartments and facts.
      */
     range?: PartialRecompRange;
+    /** @internal Exercises the lease/marker race without delaying production callers. */
+    onLeaseAcquired?: () => void;
 }
 
 export interface ExecuteContextRecompResult {
@@ -247,6 +249,19 @@ export async function executeContextRecompWithResult(
             published: false,
         };
     }
+    options.onLeaseAcquired?.();
+    if (isWrapupInProgress(deps.db, sessionId)) {
+        // Close the marker/lease race: wrapup publishes its durable ownership
+        // marker before it waits for the compartment lease, so a recomp that
+        // checked early can still win the lease after wrapup has started.
+        sessionLog(sessionId, "recomp skipped: /ctx-wrapup became active");
+        releaseCompartmentLease(deps.db, sessionId, holderId);
+        return {
+            message:
+                "## Magic Recomp — Skipped\n\n/ctx-wrapup is already compacting this session. Wait for it to finish, then try `/ctx-recomp` again.",
+            published: false,
+        };
+    }
     const renewal = startLeaseRenewal(deps, holderId);
     const runnerDeps = withPublishedCallback({ ...deps, compartmentLeaseHolderId: holderId });
     const promise = options.range

packages/plugin/src/hooks/magic-context/protected-tail-boundary.test.ts

@@ -510,6 +510,39 @@ it("fingerprints and true-raw tokens change when nested tool output grows with t
     expect(longTokens).toBeGreaterThan(shortTokens);
 });
 
+it("fingerprints same-length edits to counted content fields", () => {
+    const before: RawMessage[] = [
+        {
+            ordinal: 1,
+            id: "m1",
+            role: "assistant",
+            version: 1,
+            parts: [
+                { type: "text", text: "alpha" },
+                { type: "thinking", thinking: "bravo" },
+                { type: "tool", callID: "call-1", state: { input: { q: "one" }, output: "delta" } },
+            ],
+        },
+    ];
+    const after: RawMessage[] = [
+        {
+            ordinal: 1,
+            id: "m1",
+            role: "assistant",
+            version: 1,
+            parts: [
+                { type: "text", text: "omega" },
+                { type: "thinking", thinking: "gamma" },
+                { type: "tool", callID: "call-1", state: { input: { q: "two" }, output: "sigma" } },
+            ],
+        },
+    ];
+
+    expect(computeRawRangeFingerprint(before, 1, 2)).not.toBe(
+        computeRawRangeFingerprint(after, 1, 2),
+    );
+});
+
 it("moves a candidate boundary forward to the first later open tool invocation", () => {
     expect(
         fenceBoundaryForToolArcs(10, [{ callId: "open", invOrdinal: 20, resOrdinal: null }], 9, 10),

packages/plugin/src/hooks/magic-context/read-session-chunk.test.ts

@@ -421,5 +421,63 @@ describe("readSessionChunk", () => {
             expect(chunk.messageCount).toBe(1);
             expect(chunk.hasMore).toBe(false);
         });
+
+        it("reports hasMore false when the remaining eligible tail is only filtered noise", () => {
+            //#given
+            useTempDataHome("read-session-noise-tail-");
+            createOpenCodeDbWithMessages("ses-noise-tail", [
+                { id: "m-1", role: "user", part: { type: "text", text: "eligible work" } },
+                { id: "m-2", role: "assistant", part: { type: "text", text: "done" } },
+                {
+                    id: "m-3",
+                    role: "user",
+                    part: { type: "text", text: "## Magic Status", ignored: true },
+                },
+                {
+                    id: "m-4",
+                    role: "user",
+                    part: {
+                        type: "text",
+                        text: "<system-reminder>background finished</system-reminder>",
+                    },
+                },
+            ]);
+
+            //#when
+            const chunk = readSessionChunk("ses-noise-tail", 100_000, 1);
+
+            //#then: filtered tail noise was scanned, so no semantic work remains.
+            expect(chunk.endIndex).toBe(2);
+            expect(chunk.text).toContain("eligible work");
+            expect(chunk.text).not.toContain("Magic Status");
+            expect(chunk.hasMore).toBe(false);
+        });
+
+        it("keeps hasMore true when budget-blocked content precedes trailing noise", () => {
+            //#given
+            useTempDataHome("read-session-blocked-before-noise-");
+            createOpenCodeDbWithMessages("ses-blocked-before-noise", [
+                { id: "m-1", role: "user", part: { type: "text", text: "first content" } },
+                {
+                    id: "m-2",
+                    role: "assistant",
+                    part: { type: "text", text: "second content beyond budget" },
+                },
+                {
+                    id: "m-3",
+                    role: "user",
+                    part: { type: "text", text: "## Magic Status", ignored: true },
+                },
+            ]);
+
+            //#when
+            const chunk = readSessionChunk("ses-blocked-before-noise", 1, 1);
+
+            //#then
+            expect(chunk.endIndex).toBe(1);
+            expect(chunk.text).toContain("first content");
+            expect(chunk.text).not.toContain("second content");
+            expect(chunk.hasMore).toBe(true);
+        });
     });
 });