mason: enforce mc:protected regions after maintain-docs dreamer task

Author: ualtinok

packages/plugin/src/features/magic-context/dreamer/dreamer.test.ts

@@ -1,6 +1,9 @@
 /// <reference types="bun-types" />
 
 import { afterAll, afterEach, describe, expect, it, mock, spyOn } from "bun:test";
+import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
 import type { PluginContext } from "../../../plugin/types";
 import * as shared from "../../../shared";
 import { Database } from "../../../shared/sqlite";
@@ -275,6 +278,75 @@ describe("dreamer", () => {
             ]);
         });
 
+        it("maintain-docs restores protected regions after dreamer mutates ARCHITECTURE.md", async () => {
+            db = createTestDb();
+            const docsDir = mkdtempSync(join(tmpdir(), "mc-maintain-docs-"));
+            const startMarker = "<!-- mc:protected START — hand-authored cache-stability core. -->";
+            const endMarker = "<!-- mc:protected END -->";
+            const protectedBody = "invariant body bytes";
+            const originalArch = [
+                "# Architecture",
+                "",
+                "outside before",
+                startMarker,
+                protectedBody,
+                endMarker,
+                "outside after",
+            ].join("\n");
+            writeFileSync(join(docsDir, "ARCHITECTURE.md"), originalArch, "utf8");
+
+            const client = createDreamClient();
+            const promptSyncSpy = spyOn(
+                shared,
+                "promptSyncWithModelSuggestionRetry",
+            ).mockImplementation(async () => {
+                const tampered = [
+                    "# Architecture",
+                    "",
+                    "outside before edited",
+                    startMarker,
+                    "TAMPERED protected bytes",
+                    endMarker,
+                    "outside after edited",
+                ].join("\n");
+                writeFileSync(join(docsDir, "ARCHITECTURE.md"), tampered, "utf8");
+            });
+
+            try {
+                const result = await runDream({
+                    db,
+                    client,
+                    projectIdentity: "/repo/project",
+                    tasks: ["maintain-docs"],
+                    taskTimeoutMinutes: 5,
+                    maxRuntimeMinutes: 10,
+                    parentSessionId: "parent-1",
+                    sessionDirectory: docsDir,
+                });
+
+                expect(result.tasks.map((t) => t.name)).toEqual(["maintain-docs"]);
+                expect(result.tasks[0]?.error).toBeUndefined();
+
+                const onDisk = readFileSync(join(docsDir, "ARCHITECTURE.md"), "utf8");
+                expect(onDisk).toContain("outside before edited");
+                expect(onDisk).toContain("outside after edited");
+                expect(onDisk).toContain(protectedBody);
+                expect(onDisk).not.toContain("TAMPERED protected bytes");
+                const origBlock = originalArch.slice(
+                    originalArch.indexOf(startMarker),
+                    originalArch.indexOf(endMarker) + endMarker.length,
+                );
+                const diskBlock = onDisk.slice(
+                    onDisk.indexOf(startMarker),
+                    onDisk.indexOf(endMarker) + endMarker.length,
+                );
+                expect(diskBlock).toBe(origBlock);
+            } finally {
+                promptSyncSpy.mockRestore();
+                rmSync(docsDir, { recursive: true, force: true });
+            }
+        });
+
         it("trips circuit breaker after three consecutive identical model failures", async () => {
             db = createTestDb();
             const createdSessionIds: string[] = [];

packages/plugin/src/features/magic-context/dreamer/index.ts

@@ -1,4 +1,5 @@
 export * from "./lease";
+export * from "./protected-regions";
 export * from "./queue";
 export * from "./runner";
 export * from "./scheduler";

packages/plugin/src/features/magic-context/dreamer/maintain-docs-protected-enforcement.ts

@@ -0,0 +1,52 @@
+import { existsSync, readFileSync, writeFileSync } from "node:fs";
+import { join } from "node:path";
+import { log } from "../../../shared/logger";
+import { enforceProtectedRegions } from "./protected-regions";
+
+export const MAINTAIN_DOCS_SNAPSHOT_FILES = ["ARCHITECTURE.md", "STRUCTURE.md"] as const;
+
+export type MaintainDocsDocSnapshot = Map<string, string>;
+
+/** Read canonical pre-task bytes for maintain-docs enforcement. */
+export function snapshotMaintainDocsFiles(docsDir: string): MaintainDocsDocSnapshot {
+    const snapshot = new Map<string, string>();
+    for (const name of MAINTAIN_DOCS_SNAPSHOT_FILES) {
+        const path = join(docsDir, name);
+        try {
+            if (existsSync(path)) {
+                snapshot.set(name, readFileSync(path, "utf8"));
+            }
+        } catch {
+            // best-effort snapshot
+        }
+    }
+    return snapshot;
+}
+
+/**
+ * After maintain-docs, re-read on-disk docs and restore protected regions from the pre-task snapshot.
+ * Best-effort: read/write failures are logged, not thrown.
+ */
+export function enforceMaintainDocsProtectedRegions(args: {
+    docsDir: string;
+    snapshot: MaintainDocsDocSnapshot;
+}): void {
+    for (const [fileName, original] of args.snapshot) {
+        const path = join(args.docsDir, fileName);
+        try {
+            const current = readFileSync(path, "utf8");
+            const { text, violated } = enforceProtectedRegions(original, current);
+            if (!violated) {
+                continue;
+            }
+            writeFileSync(path, text, "utf8");
+            log(
+                `[dreamer] maintain-docs altered a protected region in ${fileName} — restored from pre-task snapshot`,
+            );
+        } catch (error) {
+            log(
+                `[dreamer] maintain-docs protected-region enforcement failed for ${fileName}: ${error}`,
+            );
+        }
+    }
+}

packages/plugin/src/features/magic-context/dreamer/protected-regions.test.ts

@@ -0,0 +1,106 @@
+/// <reference types="bun-types" />
+
+import { describe, expect, it } from "bun:test";
+import { enforceProtectedRegions, extractProtectedBlocks } from "./protected-regions";
+
+const START = "<!-- mc:protected START — hand-authored cache-stability core. Only humans edit. -->";
+const END = "<!-- mc:protected END -->";
+
+function doc(outsideBefore: string, protectedBody: string, outsideAfter: string): string {
+    return `${outsideBefore}\n${START}\n${protectedBody}\n${END}\n${outsideAfter}`;
+}
+
+describe("extractProtectedBlocks", () => {
+    it("returns empty when no markers", () => {
+        expect(extractProtectedBlocks("# Hello\n\nWorld")).toEqual([]);
+    });
+});
+
+describe("enforceProtectedRegions", () => {
+    it("identical protected region → no change, violated:false", () => {
+        const original = doc("intro", "cache core", "outro");
+        const candidate = doc("intro", "cache core", "outro");
+        const result = enforceProtectedRegions(original, candidate);
+        expect(result).toEqual({ text: candidate, violated: false });
+    });
+
+    it("edited text outside the region + untouched region → keeps edits, violated:false", () => {
+        const original = doc("intro", "cache core", "outro");
+        const candidate = doc("intro v2", "cache core", "outro v2");
+        const result = enforceProtectedRegions(original, candidate);
+        expect(result.violated).toBe(false);
+        expect(result.text).toBe(candidate);
+        expect(result.text).toContain("intro v2");
+        expect(result.text).toContain("outro v2");
+    });
+
+    it("altered protected region → restored byte-identical, violated:true, surrounding edits preserved", () => {
+        const original = doc("intro", "cache core", "outro");
+        const candidate = doc("intro v2", "cache CORE", "outro v2");
+        const result = enforceProtectedRegions(original, candidate);
+        expect(result.violated).toBe(true);
+        expect(result.text).toContain("intro v2");
+        expect(result.text).toContain("outro v2");
+        expect(result.text).toContain("cache core");
+        expect(result.text).not.toContain("cache CORE");
+        const origBlock = extractProtectedBlocks(original)[0]?.block;
+        const resultBlock = extractProtectedBlocks(result.text)[0]?.block;
+        expect(resultBlock).toBe(origBlock);
+    });
+
+    it("deleted protected markers → whole candidate rejected, returns original, violated:true", () => {
+        const original = doc("intro", "cache core", "outro");
+        const candidate = "intro v2\nno markers\noutro v2";
+        const result = enforceProtectedRegions(original, candidate);
+        expect(result).toEqual({ text: original, violated: true });
+    });
+
+    it("no protected region in original → candidate passed through unchanged", () => {
+        const original = "# STRUCTURE\n\nNo protected blocks here.";
+        const candidate = "# STRUCTURE\n\nUpdated tree.";
+        const result = enforceProtectedRegions(original, candidate);
+        expect(result).toEqual({ text: candidate, violated: false });
+    });
+
+    it("multiple protected blocks → each matched by start-marker identity and independently enforced", () => {
+        const startA = "<!-- mc:protected START block A -->";
+        const startB = "<!-- mc:protected START block B -->";
+        const original = ["head", startA, "body A", END, "mid", startB, "body B", END, "tail"].join(
+            "\n",
+        );
+
+        const candidate = [
+            "head edited",
+            startA,
+            "body A TAMPERED",
+            END,
+            "mid edited",
+            startB,
+            "body B",
+            END,
+            "tail edited",
+        ].join("\n");
+
+        const result = enforceProtectedRegions(original, candidate);
+        expect(result.violated).toBe(true);
+        expect(result.text).toContain("head edited");
+        expect(result.text).toContain("mid edited");
+        expect(result.text).toContain("body A\n");
+        expect(result.text).not.toContain("body A TAMPERED");
+        expect(extractProtectedBlocks(result.text)[0]?.block).toBe(
+            extractProtectedBlocks(original)[0]?.block,
+        );
+        expect(extractProtectedBlocks(result.text)[1]?.block).toBe(
+            extractProtectedBlocks(original)[1]?.block,
+        );
+    });
+
+    it("missing one of multiple blocks → rejects whole candidate", () => {
+        const startA = "<!-- mc:protected START block A -->";
+        const startB = "<!-- mc:protected START block B -->";
+        const original = [startA, "a", END, "x", startB, "b", END].join("\n");
+        const candidate = [startA, "a", END, "x only one block"].join("\n");
+        const result = enforceProtectedRegions(original, candidate);
+        expect(result).toEqual({ text: original, violated: true });
+    });
+});

packages/plugin/src/features/magic-context/dreamer/protected-regions.ts

@@ -0,0 +1,120 @@
+export interface EnforceProtectedRegionsResult {
+    /** The text to persist (candidate, repaired candidate, or original on reject). */
+    text: string;
+    violated: boolean;
+}
+
+const PROTECTED_START_TOKEN = "mc:protected START";
+const PROTECTED_END_TOKEN = "mc:protected END";
+
+export interface ProtectedBlock {
+    /** Full identifying start-marker line (the line containing mc:protected START). */
+    startMarkerLine: string;
+    /** Bytes from START line through END line inclusive. */
+    block: string;
+}
+
+/** Extract every mc:protected region from `text`, keyed by the full START marker line. */
+export function extractProtectedBlocks(text: string): ProtectedBlock[] {
+    const lines = text.split("\n");
+    const blocks: ProtectedBlock[] = [];
+    let i = 0;
+    while (i < lines.length) {
+        const line = lines[i];
+        if (line.includes(PROTECTED_START_TOKEN)) {
+            const startMarkerLine = line;
+            const startIdx = i;
+            while (i < lines.length && !lines[i].includes(PROTECTED_END_TOKEN)) {
+                i += 1;
+            }
+            if (i >= lines.length) {
+                break;
+            }
+            const endIdx = i;
+            const block = lines.slice(startIdx, endIdx + 1).join("\n");
+            blocks.push({ startMarkerLine, block });
+            i += 1;
+            continue;
+        }
+        i += 1;
+    }
+    return blocks;
+}
+
+function findCandidateBlockSpan(
+    candidate: string,
+    startMarkerLine: string,
+): { start: number; end: number; block: string } | null {
+    const lines = candidate.split("\n");
+    for (let i = 0; i < lines.length; i++) {
+        if (lines[i] !== startMarkerLine) {
+            continue;
+        }
+        const startIdx = i;
+        while (i < lines.length && !lines[i].includes(PROTECTED_END_TOKEN)) {
+            i += 1;
+        }
+        if (i >= lines.length) {
+            return null;
+        }
+        const endIdx = i;
+        const block = lines.slice(startIdx, endIdx + 1).join("\n");
+        return { start: startIdx, end: endIdx, block };
+    }
+    return null;
+}
+
+function spliceProtectedBlock(
+    text: string,
+    startMarkerLine: string,
+    replacementBlock: string,
+): string {
+    const lines = text.split("\n");
+    for (let i = 0; i < lines.length; i++) {
+        if (lines[i] !== startMarkerLine) {
+            continue;
+        }
+        const startIdx = i;
+        while (i < lines.length && !lines[i].includes(PROTECTED_END_TOKEN)) {
+            i += 1;
+        }
+        if (i >= lines.length) {
+            return text;
+        }
+        const endIdx = i;
+        const replacementLines = replacementBlock.split("\n");
+        const next = [...lines.slice(0, startIdx), ...replacementLines, ...lines.slice(endIdx + 1)];
+        return next.join("\n");
+    }
+    return text;
+}
+
+/**
+ * Enforce that every mc:protected region present in `original` is byte-identical
+ * in `candidate`. Returns the text to actually write and whether a violation was repaired.
+ */
+export function enforceProtectedRegions(
+    original: string,
+    candidate: string,
+): EnforceProtectedRegionsResult {
+    const originalBlocks = extractProtectedBlocks(original);
+    if (originalBlocks.length === 0) {
+        return { text: candidate, violated: false };
+    }
+
+    let text = candidate;
+    let violated = false;
+
+    for (const { startMarkerLine, block: originalBlock } of originalBlocks) {
+        const span = findCandidateBlockSpan(text, startMarkerLine);
+        if (!span) {
+            return { text: original, violated: true };
+        }
+        if (span.block !== originalBlock) {
+            text = spliceProtectedBlock(text, startMarkerLine, originalBlock);
+            violated = true;
+        }
+    }
+
+    return { text, violated };
+}