fix: gate 4 cache-busting code paths behind isCacheBustingPass

ualtinok · ualtinok · commit 34d95dc9f6ab · 2026-03-26T18:49:25.000+01:00
Finding 1 (HIGH): Compartment injection now uses in-memory cache on defer
passes. New compartments from background historian are only picked up on
execute/flush passes, preventing message[0] rewrites on cache-safe passes.

Finding 2 (HIGH): nudgePlacements.clear() after pending ops and flushed
statuses now gated behind isCacheBustingPass, preventing the two-pass
pattern where execute clears anchor then defer re-anchors.

Finding 4 (MEDIUM): When nudger() returns null on defer, existing anchor
is replayed via reinjectNudgeAtAnchor instead of cleared. Anchor
retirement only happens on cache-busting passes.

Finding 5 (MEDIUM): stripSystemInjectedMessages gated behind
isCacheBustingPass so dynamic protected-tail shifts don't remove
previously-cached messages on defer passes.

Also: early scheduler/usage computation in transform.ts eliminates
duplicate loadContextUsage/resolveSchedulerDecision calls.
diff --git a/scripts/context-dump/run-context-dump.ts b/scripts/context-dump/run-context-dump.ts
@@ -147,7 +147,7 @@ export async function runContextDump(sessionId: string): Promise<ContextDumpResu
 	try {
 		const preparedCompartmentInjection = isSubagent
 			? null
-			: prepareCompartmentInjection(contextDb, sessionId, transformedMessages as unknown as MessageLike[])
+			: prepareCompartmentInjection(contextDb, sessionId, transformedMessages as unknown as MessageLike[], true)
 		const compartmentInjection = preparedCompartmentInjection
 			? renderCompartmentInjection(sessionId, transformedMessages as unknown as MessageLike[], preparedCompartmentInjection)
 			: { injected: false, compartmentEndMessage: -1, compartmentCount: 0, skippedVisibleMessages: 0 }
diff --git a/src/hooks/magic-context/inject-compartments.ts b/src/hooks/magic-context/inject-compartments.ts
@@ -14,12 +14,25 @@ import type { MessageLike } from "./tag-messages";
 export interface PreparedCompartmentInjection {
     block: string;
     compartmentEndMessage: number;
+    compartmentEndMessageId: string;
     compartmentCount: number;
     skippedVisibleMessages: number;
     factCount: number;
     memoryCount: number;
 }
 
+/**
+ * In-memory cache of the last compartment injection result per session.
+ * On defer (cache-safe) passes, the cached result is replayed so that historian
+ * publications between passes do not bust the Anthropic prompt-cache prefix.
+ * The cache is refreshed only on cache-busting passes (execute / explicit flush).
+ */
+const injectionCache = new Map<string, PreparedCompartmentInjection>();
+
+export function clearInjectionCache(sessionId: string): void {
+    injectionCache.delete(sessionId);
+}
+
 export interface CompartmentInjectionResult {
     injected: boolean;
     compartmentEndMessage: number;
@@ -104,11 +117,30 @@ export function prepareCompartmentInjection(
     db: Database,
     sessionId: string,
     messages: MessageLike[],
+    isCacheBusting: boolean,
     projectPath?: string,
     injectionBudgetTokens?: number,
 ): PreparedCompartmentInjection | null {
+    // On defer (cache-safe) passes, replay the cached injection result so that
+    // historian publications between passes do not bust the prompt-cache prefix.
+    const cached = injectionCache.get(sessionId);
+    if (!isCacheBusting && cached) {
+        // Re-do the splice with the cached boundary (messages are rebuilt fresh each pass)
+        if (cached.compartmentEndMessageId.length > 0) {
+            const cutoffIndex = messages.findIndex(
+                (message) => message.info.id === cached.compartmentEndMessageId,
+            );
+            if (cutoffIndex >= 0) {
+                const remaining = messages.slice(cutoffIndex + 1);
+                messages.splice(0, messages.length, ...remaining);
+            }
+        }
+        return cached;
+    }
+
     const compartments = getCompartments(db, sessionId);
     if (compartments.length === 0) {
+        injectionCache.delete(sessionId);
         return null;
     }
 
@@ -122,15 +154,15 @@ export function prepareCompartmentInjection(
         // Audit note: `as` cast is safe here — session_meta schema is owned by this plugin and the two
         // columns are guaranteed present after initializeDatabase(). A type guard would add overhead on a
         // hot path (every transform) for a table we fully control.
-        const cached = db
+        const cachedMemory = db
             .prepare(
                 "SELECT memory_block_cache, memory_block_count FROM session_meta WHERE session_id = ?",
             )
             .get(sessionId) as { memory_block_cache: string; memory_block_count: number } | null;
 
-        if (cached?.memory_block_cache) {
-            memoryBlock = cached.memory_block_cache;
-            memoryCount = cached.memory_block_count;
+        if (cachedMemory?.memory_block_cache) {
+            memoryBlock = cachedMemory.memory_block_cache;
+            memoryCount = cachedMemory.memory_block_count;
         } else {
             let memories = getMemoriesByProject(db, projectPath, ["active", "permanent"]);
             if (injectionBudgetTokens && memories.length > 0) {
@@ -160,14 +192,17 @@ export function prepareCompartmentInjection(
                 compartmentEndMessage: lastEnd,
             },
         );
-        return {
+        const result: PreparedCompartmentInjection = {
             block,
             compartmentEndMessage: lastEnd,
+            compartmentEndMessageId: "",
             compartmentCount: compartments.length,
             skippedVisibleMessages: 0,
             factCount: facts.length,
             memoryCount,
         };
+        injectionCache.set(sessionId, result);
+        return result;
     }
 
     let skippedVisibleMessages = 0;
@@ -178,14 +213,17 @@ export function prepareCompartmentInjection(
         messages.splice(0, messages.length, ...remaining);
     }
 
-    return {
+    const result: PreparedCompartmentInjection = {
         block,
         compartmentEndMessage: lastEnd,
+        compartmentEndMessageId: lastEndMessageId,
         compartmentCount: compartments.length,
         skippedVisibleMessages,
         factCount: facts.length,
         memoryCount,
     };
+    injectionCache.set(sessionId, result);
+    return result;
 }
 
 export function renderCompartmentInjection(
diff --git a/src/hooks/magic-context/transform-compartment-phase.ts b/src/hooks/magic-context/transform-compartment-phase.ts
@@ -78,6 +78,7 @@ export async function runCompartmentPhase(args: RunCompartmentPhaseArgs): Promis
             args.db,
             args.resolvedSessionId,
             args.messages,
+            args.cacheAlreadyBusting ?? false,
             args.projectPath,
             args.injectionBudgetTokens,
         );
diff --git a/src/hooks/magic-context/transform-nudge-cache.test.ts b/src/hooks/magic-context/transform-nudge-cache.test.ts
@@ -128,10 +128,10 @@ describe("createTransform nudge cache handling", () => {
         expect(getPersistedNudgePlacement(db, "ses-1")).toBeNull();
     });
 
-    it("clears stored nudge placement when flushed statuses mutate content", async () => {
-        //#given
+    it("clears stored nudge placement when flushed statuses mutate content on execute pass", async () => {
+        //#given — execute pass so nudge placement clear is allowed
         useTempDataHome("context-transform-nudge-flushed-");
-        const scheduler: Scheduler = { shouldExecute: mock(() => "defer" as const) };
+        const scheduler: Scheduler = { shouldExecute: mock(() => "execute" as const) };
         const db = openDatabase();
         const nudgePlacements = createNudgePlacementStore(db);
         const contextUsageMap = new Map<string, { usage: ContextUsage; updatedAt: number }>([
diff --git a/src/hooks/magic-context/transform-postprocess-phase.ts b/src/hooks/magic-context/transform-postprocess-phase.ts
@@ -19,6 +19,7 @@ import {
     renderCompartmentInjection,
 } from "./inject-compartments";
 import { getNoteNudgeText } from "./note-nudger";
+import { reinjectNudgeAtAnchor } from "./nudge-injection";
 import type { NudgePlacementStore } from "./nudge-placement-store";
 import type { ContextNudge } from "./nudger";
 import {
@@ -101,6 +102,8 @@ export function runPostTransformPhase(args: RunPostTransformPhaseArgs): void {
     const hasPendingUserOps = pendingOps.length > 0;
     const shouldApplyPendingOps =
         (args.schedulerDecision === "execute" || isExplicitFlush) && !compartmentRunning;
+    // Central cache-busting gate used by all mutation paths below.
+    const isCacheBustingPass = isExplicitFlush || shouldApplyPendingOps;
     const shouldRunHeuristics =
         args.fullFeatureMode &&
         !compartmentRunning &&
@@ -228,7 +231,9 @@ export function runPostTransformPhase(args: RunPostTransformPhaseArgs): void {
         updateSessionMeta(args.db, args.sessionId, { lastTransformError: getErrorMessage(error) });
         args.nudgePlacements.clear(args.sessionId);
     }
-    if (didMutateFromPendingOperations) {
+    // Only clear nudge placements on cache-busting passes. Clearing on defer would
+    // cause the next pass to re-anchor the nudge on a cached assistant message (Finding 2).
+    if (didMutateFromPendingOperations && isCacheBustingPass) {
         args.nudgePlacements.clear(args.sessionId);
     }
 
@@ -274,25 +279,28 @@ export function runPostTransformPhase(args: RunPostTransformPhaseArgs): void {
     }
 
     // Remove system-injected messages (notifications, reminders, internal markers)
-    // that are OUTSIDE the protected tail. Recent ones in the tail may contain
-    // actionable info (e.g., background task IDs the agent needs for background_output).
-    // Use the configured protected_tags count to find the boundary — messages in the
-    // last N tags are kept.
-    const protectedTailStart = Math.max(0, args.messages.length - args.protectedTags * 2);
-    const strippedSystemInjected = stripSystemInjectedMessages(args.messages, protectedTailStart);
-    if (strippedSystemInjected > 0) {
-        sessionLog(
-            args.sessionId,
-            `stripped ${strippedSystemInjected} system-injected messages (notifications/reminders)`,
+    // that are OUTSIDE the protected tail. Only strip on cache-busting passes because
+    // the dynamic protected tail boundary can shift as conversation grows, which would
+    // remove previously-cached messages on defer passes (Finding 5).
+    if (isCacheBustingPass) {
+        const protectedTailStart = Math.max(0, args.messages.length - args.protectedTags * 2);
+        const strippedSystemInjected = stripSystemInjectedMessages(
+            args.messages,
+            protectedTailStart,
         );
+        if (strippedSystemInjected > 0) {
+            sessionLog(
+                args.sessionId,
+                `stripped ${strippedSystemInjected} system-injected messages (notifications/reminders)`,
+            );
+        }
     }
 
     const pendingUserTurnReminder = getPersistedStickyTurnReminder(args.db, args.sessionId);
     if (pendingUserTurnReminder) {
         // Only clear the reminder when the pass is already cache-busting (execute/flush).
         // Clearing on a cache-safe pass would remove text from an anchored user message,
         // changing cached content and busting the Anthropic prompt-cache prefix.
-        const isCacheBustingPass = isExplicitFlush || shouldApplyPendingOps || shouldRunHeuristics;
         if (args.hasRecentReduceCall && isCacheBustingPass) {
             clearPersistedStickyTurnReminder(args.db, args.sessionId);
             sessionLog(
@@ -351,8 +359,22 @@ export function runPostTransformPhase(args: RunPostTransformPhaseArgs): void {
             const t9 = performance.now();
             applyContextNudge(args.messages, nudge, args.nudgePlacements, args.sessionId);
             logTransformTiming(args.sessionId, "applyContextNudge", t9);
-        } else {
+        } else if (isCacheBustingPass) {
+            // Only retire the nudge anchor on cache-busting passes (Finding 4).
+            // Clearing on defer would remove previously-injected nudge text from
+            // the cached assistant message.
             args.nudgePlacements.clear(args.sessionId);
+        } else {
+            // Defer pass: replay existing anchor to keep cached content stable.
+            const existing = args.nudgePlacements.get(args.sessionId);
+            if (existing) {
+                reinjectNudgeAtAnchor(
+                    args.messages,
+                    existing.nudgeText,
+                    args.nudgePlacements,
+                    args.sessionId,
+                );
+            }
         }
 
         const deferredNoteText = getNoteNudgeText(args.db, args.sessionId);
diff --git a/src/hooks/magic-context/transform.ts b/src/hooks/magic-context/transform.ts
@@ -105,6 +105,18 @@ export function createTransform(deps: TransformDeps) {
         const canRunCompartments =
             fullFeatureMode && deps.client !== undefined && compartmentDirectory.length > 0;
 
+        // Compute cache-busting status early so compartment injection can use it.
+        // The scheduler and flush state are available before tagging.
+        const contextUsageEarly = loadContextUsage(deps.contextUsageMap, db, sessionId);
+        const schedulerDecisionEarly = resolveSchedulerDecision(
+            deps.scheduler,
+            sessionMeta,
+            contextUsageEarly,
+            sessionId,
+        );
+        const isCacheBusting =
+            deps.flushedSessions.has(sessionId) || schedulerDecisionEarly === "execute";
+
         let pendingCompartmentInjection: PreparedCompartmentInjection | null = null;
         if (fullFeatureMode) {
             const projectPath = deps.memoryConfig?.enabled
@@ -114,6 +126,7 @@ export function createTransform(deps: TransformDeps) {
                 db,
                 sessionId,
                 messages,
+                isCacheBusting,
                 projectPath,
                 deps.memoryConfig?.injectionBudgetTokens,
             );
@@ -165,9 +178,10 @@ export function createTransform(deps: TransformDeps) {
             logTransformTiming(sessionId, "batchFinalize:flushed", t2);
         } catch (error) {
             sessionLog(sessionId, "transform failed applying flushed statuses:", error);
-            deps.nudgePlacements.clear(sessionId);
+            // Only clear on cache-busting passes to avoid re-anchor on next defer (Finding 2).
+            if (isCacheBusting) deps.nudgePlacements.clear(sessionId);
         }
-        if (didMutateFromFlushedStatuses) {
+        if (didMutateFromFlushedStatuses && isCacheBusting) {
             deps.nudgePlacements.clear(sessionId);
         }
 
@@ -196,15 +210,9 @@ export function createTransform(deps: TransformDeps) {
             }
         }
 
-        const contextUsage = loadContextUsage(deps.contextUsageMap, db, sessionId);
-        const schedulerDecision = resolveSchedulerDecision(
-            deps.scheduler,
-            sessionMeta,
-            contextUsage,
-            sessionId,
-        );
-        const isCacheBusting =
-            deps.flushedSessions.has(sessionId) || schedulerDecision === "execute";
+        // Reuse the early scheduler result — inputs haven't changed.
+        const contextUsage = contextUsageEarly;
+        const schedulerDecision = schedulerDecisionEarly;
         const rawGetNotifParams = deps.getNotificationParams;
         const compartmentPhase = await runCompartmentPhase({
             canRunCompartments,
