@retaw106/mmc-oc 2.4.0 → 2.6.0

package/src/shared/system-directive.test.ts

@@ -0,0 +1,37 @@
+import { describe, expect, test } from "bun:test";
+import { removeNestedBlocks, removeSystemReminders } from "./system-directive";
+
+describe("removeSystemReminders", () => {
+    test("removes nested reminders without leaking the outer tail", () => {
+        expect(
+            removeSystemReminders(
+                "before <system-reminder>outer <system-reminder>inner</system-reminder> tail</system-reminder> after",
+            ),
+        ).toBe("before  after");
+    });
+
+    test("removes unclosed reminders through the end", () => {
+        expect(
+            removeSystemReminders(
+                " keep <system-reminder>hidden <system-reminder>nested</system-reminder> tail",
+            ),
+        ).toBe("keep");
+    });
+
+    test("handles attributes, case, self-closing tags and adjacent blocks", () => {
+        expect(
+            removeSystemReminders(
+                'a<SYSTEM-REMINDER id="x">hidden</SYSTEM-REMINDER><system-reminder/>b<system-reminder>hidden</system-reminder>c',
+            ),
+        ).toBe("abc");
+    });
+
+    test("preserves ordinary content and other tags", () => {
+        expect(removeSystemReminders(" <other>visible</other> ")).toBe("<other>visible</other>");
+    });
+});
+
+test("removeNestedBlocks tracks multiple tag names and only closes the matching top", () => {
+    expect(removeNestedBlocks("a<one>x<two>y</one>z</two>tail</one>b", ["one", "two"])).toBe("ab");
+    expect(removeNestedBlocks(" a<one/>b ", ["one"])).toBe(" ab ");
+});

package/src/shared/system-directive.ts

@@ -5,5 +5,23 @@ export function isSystemDirective(text: string): boolean {
 }
 
 export function removeSystemReminders(text: string): string {
-    return text.replace(/<system-reminder>[\s\S]*?<\/system-reminder>/gi, "").trim();
+    return removeNestedBlocks(text, ["system-reminder"]).trim();
+}
+
+/** Remove matching tagged regions, including nested regions and regions with no closing tag. */
+export function removeNestedBlocks(text: string, names: readonly string[]): string {
+    const tags = new RegExp(`<(/?)(${names.join("|")})(?:\\s[^<>]*?)?\\s*(/?)>`, "gi");
+    const stack: string[] = [];
+    let end = 0;
+    let result = "";
+    for (const match of text.matchAll(tags)) {
+        if (stack.length === 0) result += text.slice(end, match.index);
+        const name = match[2].toLowerCase();
+        if (match[1]) {
+            if (stack[stack.length - 1] === name) stack.pop();
+        } else if (!match[3]) stack.push(name);
+        end = match.index + match[0].length;
+    }
+    if (stack.length === 0) result += text.slice(end);
+    return result;
 }

package/src/shared/tag-transcript.ts

@@ -75,6 +75,8 @@ import type { Transcript, TranscriptMessage, TranscriptPart } from "./transcript
 export { TEXT_TAG_IDENTITY_MARKER };
 
 export interface TagTranscriptOptions {
+    /** Independently validated text sources; does not authorize tool accounting reuse. */
+    reuseTextMessageIds?: ReadonlySet<string>;
     /**
      * When true, skip injecting `§N§` prefix into visible text. Tags
      * still get assigned in the DB so historian/drops can reference
@@ -305,6 +307,12 @@ interface ToolAggregate {
     maxByteSize: number;
     /** Token count paired with maxByteSize (the same output occurrence). */
     maxTokenCount: number;
+    /**
+     * Whether maxTokenCount is a real count. Legacy rows can store bytes with a
+     * NULL token_count; seeding 0 from that must not look like a stored count,
+     * or unchanged results would never be recounted.
+     */
+    tokenCountKnown: boolean;
     /** Tool name from the first occurrence we see one on. */
     toolName: string | null;
     /** Input byte size from the invocation occurrence (for storage projection). */
@@ -417,7 +425,10 @@ export function tagTranscript(
                     targets,
                     skipPrefixInjection,
                     entryFingerprint: options.entryFingerprintByMessageId?.get(messageId) ?? null,
-                    reuseIdentity: reuseIdentity || contentDerivedTextIds !== undefined,
+                    reuseIdentity:
+                        reuseIdentity ||
+                        options.reuseTextMessageIds?.has(messageId) === true ||
+                        contentDerivedTextIds !== undefined,
                     timing,
                     textIdentitySourceCache: options.textIdentitySourceCache,
                     textTokenCache: options.textTokenCache,
@@ -508,6 +519,13 @@ export function tagTranscript(
                             timing,
                             options.toolTokenCache,
                             tokenCacheKey,
+                            {
+                                byteSize: existing.maxByteSize,
+                                tokenCount: existing.tokenCountKnown
+                                    ? existing.maxTokenCount
+                                    : null,
+                                inputTokenCount: existing.inputTokenCount,
+                            },
                         );
                         text = accounting.text;
                         if (part.kind === "tool_result") {
@@ -596,6 +614,7 @@ export function tagTranscript(
                         // genuinely larger folded result can still bump accounting.
                         maxByteSize: reusableAccounting.byteSize,
                         maxTokenCount: reusableAccounting.tokenCount ?? 0,
+                        tokenCountKnown: reusableAccounting.tokenCount !== null,
                         toolName: null,
                         inputByteSize: reusableAccounting.inputByteSize,
                         inputTokenCount: reusableAccounting.inputTokenCount,
@@ -623,6 +642,7 @@ export function tagTranscript(
                         timing,
                         options.toolTokenCache,
                         tokenCacheKey,
+                        tagger.getToolTagAccounting(sessionId, callId, messageId),
                     );
                     text = accounting.text;
                     const outputByteSize = part.kind === "tool_result" ? accounting.byteSize : 0;
@@ -657,6 +677,11 @@ export function tagTranscript(
                         occurrences: [{ message, part, kind: part.kind }],
                         maxByteSize: persistedAccounting?.byteSize ?? outputByteSize,
                         maxTokenCount: persistedAccounting?.tokenCount ?? outputTokenCount,
+                        // A fresh allocation just wrote outputTokenCount (0 for an
+                        // invocation); an existing row keeps its NULL as unknown.
+                        tokenCountKnown: persistedAccounting
+                            ? persistedAccounting.tokenCount !== null
+                            : true,
                         toolName: accounting.toolName,
                         inputByteSize:
                             persistedAccounting?.inputByteSize ??
@@ -745,7 +770,9 @@ function syncToolAggregateAccounting(
 ): void {
     tagger.setToolTagAccounting(sessionId, aggregate.tagId, {
         byteSize: aggregate.maxByteSize,
-        tokenCount: aggregate.maxTokenCount,
+        // Unknown stays unknown so the next occurrence/pass recounts instead of
+        // caching a fabricated 0 for a legacy row.
+        tokenCount: aggregate.tokenCountKnown ? aggregate.maxTokenCount : null,
         inputByteSize: aggregate.inputByteSize,
         inputTokenCount: aggregate.inputTokenCount,
     });
@@ -770,6 +797,7 @@ function applyGrownToolResultAccounting(args: GrownToolResultAccountingArgs): vo
     }
     args.aggregate.maxByteSize = args.byteSize;
     args.aggregate.maxTokenCount = tokenCount;
+    args.aggregate.tokenCountKnown = true;
     updateTagByteSize(args.db, args.sessionId, args.aggregate.tagId, args.byteSize);
     updateTagTokenCount(args.db, args.sessionId, args.aggregate.tagId, tokenCount);
     syncToolAggregateAccounting(args.tagger, args.sessionId, args.aggregate);
@@ -780,12 +808,22 @@ function readAggregateToolAccounting(
     timing: TagTranscriptTiming | undefined,
     tokenCache?: Map<string, { text: string; tokenCount: number }>,
     tokenCacheKey?: string,
+    persisted?: { byteSize: number; tokenCount: number | null; inputTokenCount: number | null },
 ): AggregateToolAccounting {
     const text = part.getText() ?? "";
     const byteSize = getToolPartByteSize(part, text);
-    const metadata = part.getToolMetadata();
+    // Persisted tool-result accounting tracks the largest byte size seen. Skip
+    // output tokenization when that size has not grown and a count is stored;
+    // known input counts also need no tokenization. Growth and missing counts still take
+    // the original counting path, even when the identity was not reusable.
+    const metadata = part.getToolMetadata({
+        skipInputTokenCount: persisted?.inputTokenCount != null,
+    });
     let tokenCount = 0;
-    if (part.kind === "tool_result") {
+    if (
+        part.kind === "tool_result" &&
+        (!persisted || persisted.tokenCount === null || byteSize > persisted.byteSize)
+    ) {
         const tokenStart = timing ? performance.now() : 0;
         tokenCount = getCachedToolPartTokenCount(part, text, tokenCache, tokenCacheKey);
         if (timing) timing.tokenCounting += performance.now() - tokenStart;
@@ -835,12 +873,17 @@ class ConversationEndGuard {
         }
     }
 
-    wouldStrand(occurrences: readonly ToolOccurrence[]): boolean {
+    wouldStrand(
+        occurrences: readonly ToolOccurrence[],
+        plannedRemovals?: ReadonlySet<unknown>,
+    ): boolean {
         const end = this.end;
         if (!end || !occurrences.some((occ) => occ.message === end)) return false;
         const leaving = new Set(occurrences.map((occ) => partKey(occ.part)));
         const survives = (part: TranscriptPart) =>
-            !leaving.has(partKey(part)) && !this.removed.has(partKey(part));
+            !leaving.has(partKey(part)) &&
+            !this.removed.has(partKey(part)) &&
+            !plannedRemovals?.has(partKey(part));
         return end.info.role === "assistant"
             ? !end.parts.some((part) => part.kind === "tool_result" && survives(part))
             : !end.parts.some((part) => partHasContent(part) && survives(part));
@@ -1379,8 +1422,11 @@ function buildAggregateTarget(
             }
             return null;
         },
-        wouldStrandConversationEnd(): boolean {
-            return conversationEnd?.wouldStrand(occurrences) ?? false;
+        wouldStrandConversationEnd(plannedRemovals): boolean {
+            return conversationEnd?.wouldStrand(occurrences, plannedRemovals) ?? false;
+        },
+        notePlannedRemoval(plannedRemovals): void {
+            for (const occ of occurrences) plannedRemovals.add(partKey(occ.part));
         },
         editMarker(): "truncated" | "absent" {
             // Edit-marker: preserve the tool_use input's filePath + a region
@@ -1406,6 +1452,11 @@ function buildAggregateTarget(
         },
         // Open invocations still belong to the current turn; only complete arcs reclaim.
         prepareDrop: () => !occurrences.some((occ) => occ.part.canRemove?.() === "defer"),
+        prepareReplayDrop: () => {
+            if (complete() && !requiresToolArcSkeleton) {
+                for (const occ of occurrences) occ.part.canRemove?.();
+            }
+        },
         canDrop: () => complete() && !occurrences.some((occ) => occ.part.hasUserAnswer?.()),
         requiresToolArcSkeleton,
         // Non-mutating read of the invocation input (the tool_use occurrence
@@ -1443,6 +1494,10 @@ function buildTextTarget(
     return {
         mutationParts: [{ message: { info: message.info, parts: [] }, part }],
         textPrefix,
+        canSetContent(content): boolean {
+            const current = part.getText();
+            return typeof current === "string" && current !== content;
+        },
         setContent(content: string): boolean {
             return part.setText(content);
         },

package/src/shared/transcript.ts

@@ -132,7 +132,7 @@ export interface TranscriptPart {
      *
      * For non-tool parts both byte fields are undefined/0.
      */
-    getToolMetadata(): {
+    getToolMetadata(options?: { skipInputTokenCount?: boolean }): {
         toolName: string | undefined;
         inputByteSize: number;
         inputTokenCount: number;

package/src/shared/write-transaction-attribution.test.ts

@@ -14,11 +14,7 @@ const fencedSources: Array<{ path: string; sites: string[]; reporterCalls?: stri
     },
     {
         path: "packages/plugin/src/features/magic-context/message-index.ts",
-        sites: [
-            "message_index_incremental",
-            "message_index_reconcile",
-            "message_index_orphan_sweep",
-        ],
+        sites: ["message_index_incremental", "message_index_reconcile"],
     },
     {
         path: "packages/plugin/src/features/magic-context/storage-meta-persisted.ts",
@@ -101,7 +97,6 @@ const newlyCoveredSites = [
     "opencode_soft_refresh_cache",
     "message_index_incremental",
     "message_index_reconcile",
-    "message_index_orphan_sweep",
     "message_index_clear",
     "message_fts_rowid_backfill",
     "storage_meta_wrapup_expiry_cleanup",

package/dist/features/magic-context/message-index-sweep.d.ts

@@ -1,11 +0,0 @@
-import type { Database } from "../../shared/sqlite";
-import { type MessageHistoryOrphanSweepOptions, type MessageHistoryOrphanSweepResult } from "./message-index";
-/**
- * Scope wrapper for the background sweep. It lives beside `message-index` rather
- * than inside it: importing the transform-pass scope helpers into
- * `message-index.ts` itself changes the bundled module graph in a way that hits
- * a Bun 1.4.2 packaging bug (a top-level `__promiseAll` init call is emitted
- * without its definition, which crashes node hosts at import time).
- */
-export declare function sweepOrphanedOpenCodeMessageIndexes(db: Database, openReadableOpenCodeDb: () => Database | null, options?: MessageHistoryOrphanSweepOptions): MessageHistoryOrphanSweepResult;
-//# sourceMappingURL=message-index-sweep.d.ts.map

package/dist/features/magic-context/message-index-sweep.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"message-index-sweep.d.ts","sourceRoot":"","sources":["../../../src/features/magic-context/message-index-sweep.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,QAAQ,EAAE,MAAM,qBAAqB,CAAC;AAEpD,OAAO,EACH,KAAK,gCAAgC,EACrC,KAAK,+BAA+B,EAEvC,MAAM,iBAAiB,CAAC;AAEzB;;;;;;GAMG;AACH,wBAAgB,mCAAmC,CAC/C,EAAE,EAAE,QAAQ,EACZ,sBAAsB,EAAE,MAAM,QAAQ,GAAG,IAAI,EAC7C,OAAO,GAAE,gCAAqC,GAC/C,+BAA+B,CAMjC"}