@cortexkit/opencode-magic-context 0.31.4 → 0.32.0

package/src/shared/tag-transcript.ts

@@ -83,6 +83,21 @@ export interface TagTranscriptOptions {
      * contentId) since all parts of a message share one fingerprint.
      */
     entryFingerprintByMessageId?: ReadonlyMap<string, string>;
+    /**
+     * Stable Pi message ids observed on a prior pass. Their immutable parts may
+     * reuse tag assignments while this pass still reapplies visible prefixes and
+     * rebuilds the complete set of messages affected by each tag.
+     */
+    reuseMessageIds?: ReadonlySet<string>;
+    /** Exact text/count pairs retained by Pi for safe lazy-token backfill reuse. */
+    textTokenCache?: Map<string, { text: string; tokenCount: number }>;
+    /** Exact tool-result text/count pairs retained under composite tag identity. */
+    toolTokenCache?: Map<string, { text: string; tokenCount: number }>;
+    /** Optional process-local benchmark callback; production callers omit it. */
+    onTiming?: (
+        phase: "identity" | "prefix" | "targets" | "tokenCounting",
+        elapsedMs: number,
+    ) => void;
 }
 
 export interface TagTranscriptResult {
@@ -134,8 +149,17 @@ interface ToolOccurrence {
     kind: "tool_use" | "tool_result";
 }
 
+interface TagTranscriptTiming {
+    identity: number;
+    prefix: number;
+    targets: number;
+    tokenCounting: number;
+}
+
 interface ToolAggregate {
     callId: string;
+    /** Every occurrence seen so far belongs to a previously tagged stable message. */
+    identityReusable: boolean;
     occurrences: ToolOccurrence[];
     /** Largest byteSize seen across occurrences — used as the tag size. */
     maxByteSize: number;
@@ -145,8 +169,8 @@ interface ToolAggregate {
     toolName: string | null;
     /** Input byte size from the invocation occurrence (for storage projection). */
     inputByteSize: number;
-    /** Whether input_token_count has been persisted (from the tool_use occurrence). */
-    inputTokenStored: boolean;
+    /** Persisted input-token count, or null for a legacy/unobserved invocation. */
+    inputTokenCount: number | null;
 }
 
 export function tagTranscript(
@@ -158,6 +182,9 @@ export function tagTranscript(
 ): TagTranscriptResult {
     const skipPrefixInjection = options.skipPrefixInjection === true;
     const targets = new Map<number, TagTarget>();
+    const timing: TagTranscriptTiming | undefined = options.onTiming
+        ? { identity: 0, prefix: 0, targets: 0, tokenCounting: 0 }
+        : undefined;
 
     // Tool aggregation is keyed by the same owner+callId identity used by
     // assignToolTag. OpenCode/Pi callId counters can repeat across turns, so
@@ -179,6 +206,8 @@ export function tagTranscript(
         if (message === undefined) continue;
         activeToolResultRun = undefined;
         const messageId = message.info.id;
+        const reuseIdentity =
+            messageId !== undefined && options.reuseMessageIds?.has(messageId) === true;
 
         let textOrdinal = 0;
         const parts = message.parts;
@@ -213,6 +242,9 @@ export function tagTranscript(
                     targets,
                     skipPrefixInjection,
                     entryFingerprint: options.entryFingerprintByMessageId?.get(messageId) ?? null,
+                    reuseIdentity,
+                    timing,
+                    textTokenCache: options.textTokenCache,
                 });
                 textOrdinal += 1;
                 continue;
@@ -224,13 +256,8 @@ export function tagTranscript(
                     continue;
                 }
 
+                const identityStart = timing ? performance.now() : 0;
                 const callId = part.id;
-                const text = part.getText() ?? "";
-                const toolByteSize = getToolPartByteSize(part, text);
-                const toolTokenCount = getToolPartTokenCount(part, text);
-                const meta = part.getToolMetadata();
-                const inputTokenCount = meta.inputTokenCount;
-
                 if (typeof callId !== "string" || callId.length === 0) {
                     activeToolResultRun = undefined;
                     // No stable callId to aggregate on. Tag independently.
@@ -245,6 +272,8 @@ export function tagTranscript(
                         db,
                         targets,
                         skipPrefixInjection,
+                        reuseIdentity,
+                        timing,
                     });
                     continue;
                 }
@@ -267,58 +296,92 @@ export function tagTranscript(
                 const aggregateKey: string = existingKey ?? makeToolCompositeKey(messageId, callId);
                 const existing = toolAggregates.get(aggregateKey);
                 if (existing) {
-                    // Later occurrence for this owner+callId pair. Merge into the
-                    // aggregate, update byte accounting if larger, and rebuild the
-                    // TagTarget so drops mutate both invocation and result.
-                    existing.occurrences.push({
-                        message,
+                    existing.occurrences.push({ message, part, kind: part.kind });
+                    const canReuseIdentity = reuseIdentity && existing.identityReusable;
+                    let text = "";
+                    if (canReuseIdentity) {
+                        // Prefix replay always needs result text. Byte length is also
+                        // cheap enough to guard the durable growth invariant; BPE and
+                        // DB writes remain staged behind an actual persisted-size bump.
+                        if (part.kind === "tool_result") {
+                            text = part.getText() ?? "";
+                            applyGrownToolResultAccounting({
+                                db,
+                                sessionId,
+                                tagger,
+                                aggregate: existing,
+                                byteSize: getToolPartByteSize(part, text),
+                                part,
+                                text,
+                                timing,
+                                tokenCache: options.toolTokenCache,
+                                tokenCacheKey: aggregateKey,
+                            });
+                        }
+                        if (timing) timing.identity += performance.now() - identityStart;
+                    } else {
+                        const accounting = readAggregateToolAccounting(
+                            part,
+                            timing,
+                            options.toolTokenCache,
+                            aggregateKey,
+                        );
+                        text = accounting.text;
+                        if (part.kind === "tool_result") {
+                            applyGrownToolResultAccounting({
+                                db,
+                                sessionId,
+                                tagger,
+                                aggregate: existing,
+                                byteSize: accounting.byteSize,
+                                part,
+                                text,
+                                timing,
+                                knownTokenCount: accounting.tokenCount,
+                            });
+                        }
+                        if (existing.toolName === null && accounting.toolName) {
+                            existing.toolName = accounting.toolName;
+                        }
+                        if (
+                            existing.inputByteSize === 0 &&
+                            part.kind === "tool_use" &&
+                            accounting.inputByteSize > 0
+                        ) {
+                            existing.inputByteSize = accounting.inputByteSize;
+                            updateTagInputByteSize(
+                                db,
+                                sessionId,
+                                existing.tagId,
+                                accounting.inputByteSize,
+                            );
+                        }
+                        if (
+                            existing.inputTokenCount === null &&
+                            part.kind === "tool_use" &&
+                            accounting.inputTokenCount > 0
+                        ) {
+                            existing.inputTokenCount = accounting.inputTokenCount;
+                            updateTagInputTokenCount(
+                                db,
+                                sessionId,
+                                existing.tagId,
+                                accounting.inputTokenCount,
+                            );
+                        }
+                        syncToolAggregateAccounting(tagger, sessionId, existing);
+                        if (timing) timing.identity += performance.now() - identityStart;
+                    }
+                    existing.identityReusable &&= reuseIdentity;
+                    applyToolPrefixAndTarget({
+                        skipPrefixInjection,
                         part,
-                        kind: part.kind,
+                        text,
+                        tagId: existing.tagId,
+                        aggregate: existing,
+                        targets,
+                        timing,
                     });
-                    // byte_size tracks OUTPUT bytes only (the tool_result
-                    // occurrence). The invocation args are captured separately in
-                    // inputByteSize below; counting tool_use bytes here too would
-                    // double-count args in the emergency-drop reclaim formula
-                    // (byteSize + inputByteSize + reasoning). Mirrors OpenCode,
-                    // which assigns the tool tag on the result path.
-                    if (part.kind === "tool_result" && toolByteSize > existing.maxByteSize) {
-                        existing.maxByteSize = toolByteSize;
-                        existing.maxTokenCount = toolTokenCount;
-                        updateTagByteSize(db, sessionId, existing.tagId, toolByteSize);
-                        // Keep token_count in lockstep with the byte bump so the
-                        // grown output's tokens aren't undercounted by readers.
-                        updateTagTokenCount(db, sessionId, existing.tagId, toolTokenCount);
-                    }
-                    if (existing.toolName === null && meta.toolName) {
-                        existing.toolName = meta.toolName;
-                    }
-                    if (
-                        existing.inputByteSize === 0 &&
-                        part.kind === "tool_use" &&
-                        meta.inputByteSize > 0
-                    ) {
-                        existing.inputByteSize = meta.inputByteSize;
-                        updateTagInputByteSize(db, sessionId, existing.tagId, meta.inputByteSize);
-                    }
-                    if (
-                        !existing.inputTokenStored &&
-                        part.kind === "tool_use" &&
-                        inputTokenCount > 0
-                    ) {
-                        existing.inputTokenStored = true;
-                        updateTagInputTokenCount(db, sessionId, existing.tagId, inputTokenCount);
-                    }
-                    // Inject §N§ prefix into this tool_result occurrence
-                    // (matches OpenCode behavior — only result gets the prefix).
-                    if (!skipPrefixInjection && part.kind === "tool_result") {
-                        part.setText(prependTag(existing.tagId, text));
-                    }
-                    // Rebuild the aggregate target so it walks the now-
-                    // longer occurrences list.
-                    targets.set(
-                        existing.tagId,
-                        buildAggregateTarget(existing.tagId, existing.occurrences),
-                    );
                     if (part.kind === "tool_result") {
                         markToolAggregateResolved(
                             callId,
@@ -327,18 +390,61 @@ export function tagTranscript(
                         );
                         activeToolResultRun = { callId, aggregateKey };
                     }
+                    continue;
+                }
+
+                const reusableTagId = reuseIdentity
+                    ? tagger.getToolTag(sessionId, callId, messageId)
+                    : undefined;
+                const reusableAccounting = reuseIdentity
+                    ? tagger.getToolTagAccounting(sessionId, callId, messageId)
+                    : undefined;
+                let aggregate: ToolAggregate & { tagId: number };
+                let text = "";
+                if (reusableTagId !== undefined && reusableAccounting !== undefined) {
+                    aggregate = {
+                        callId,
+                        tagId: reusableTagId,
+                        identityReusable: true,
+                        occurrences: [{ message, part, kind: part.kind }],
+                        // Stable identity does not imply stable payload size. Seed from
+                        // the persisted row so unchanged results avoid BPE while a
+                        // genuinely larger folded result can still bump accounting.
+                        maxByteSize: reusableAccounting.byteSize,
+                        maxTokenCount: reusableAccounting.tokenCount ?? 0,
+                        toolName: null,
+                        inputByteSize: reusableAccounting.inputByteSize,
+                        inputTokenCount: reusableAccounting.inputTokenCount,
+                    };
+                    if (part.kind === "tool_result") {
+                        text = part.getText() ?? "";
+                        applyGrownToolResultAccounting({
+                            db,
+                            sessionId,
+                            tagger,
+                            aggregate,
+                            byteSize: getToolPartByteSize(part, text),
+                            part,
+                            text,
+                            timing,
+                            tokenCache: options.toolTokenCache,
+                            tokenCacheKey: aggregateKey,
+                        });
+                    }
+                    if (timing) timing.identity += performance.now() - identityStart;
                 } else {
-                    // First occurrence for this owner+callId identity — reserve
-                    // the tag number. Owner stays stable across passes because
-                    // transcript message ids are durable.
-                    // byte_size is OUTPUT-only (0 until the tool_result occurrence
-                    // is seen); the invocation args live in inputByteSize. This
-                    // keeps the emergency-drop reclaim formula (byteSize +
-                    // inputByteSize + reasoning) from double-counting args when the
-                    // first occurrence is a large tool_use. Mirrors OpenCode.
-                    const outputByteSize = part.kind === "tool_result" ? toolByteSize : 0;
-                    const outputTokenCount = part.kind === "tool_result" ? toolTokenCount : 0;
-                    const firstInputTokenCount = part.kind === "tool_use" ? inputTokenCount : 0;
+                    const accounting = readAggregateToolAccounting(
+                        part,
+                        timing,
+                        options.toolTokenCache,
+                        aggregateKey,
+                    );
+                    text = accounting.text;
+                    const outputByteSize = part.kind === "tool_result" ? accounting.byteSize : 0;
+                    const outputTokenCount =
+                        part.kind === "tool_result" ? accounting.tokenCount : 0;
+                    const firstInputTokenCount =
+                        part.kind === "tool_use" ? accounting.inputTokenCount : 0;
                     const tagId = tagger.assignToolTag(
                         sessionId,
                         callId,
@@ -346,58 +452,188 @@ export function tagTranscript(
                         outputByteSize,
                         db,
                         0,
-                        meta.toolName ?? null,
-                        meta.inputByteSize,
+                        accounting.toolName,
+                        accounting.inputByteSize,
                         () => ({
                             tokenCount: outputTokenCount,
                             inputTokenCount: firstInputTokenCount,
                             reasoningTokenCount: null,
                         }),
                     );
-                    const aggregate = {
+                    const persistedAccounting = tagger.getToolTagAccounting(
+                        sessionId,
+                        callId,
+                        messageId,
+                    );
+                    aggregate = {
                         callId,
                         tagId,
-                        occurrences: [
-                            {
-                                message,
-                                part,
-                                kind: part.kind,
-                            },
-                        ],
-                        maxByteSize: outputByteSize,
-                        maxTokenCount: outputTokenCount,
-                        toolName: meta.toolName ?? null,
-                        inputByteSize: part.kind === "tool_use" ? meta.inputByteSize : 0,
-                        inputTokenStored: part.kind === "tool_use" && firstInputTokenCount > 0,
+                        identityReusable: false,
+                        occurrences: [{ message, part, kind: part.kind }],
+                        maxByteSize: persistedAccounting?.byteSize ?? outputByteSize,
+                        maxTokenCount: persistedAccounting?.tokenCount ?? outputTokenCount,
+                        toolName: accounting.toolName,
+                        inputByteSize:
+                            persistedAccounting?.inputByteSize ??
+                            (part.kind === "tool_use" ? accounting.inputByteSize : 0),
+                        inputTokenCount:
+                            persistedAccounting?.inputTokenCount ??
+                            (part.kind === "tool_use" ? firstInputTokenCount : null),
                     };
-                    toolAggregates.set(aggregateKey, aggregate);
-                    if (part.kind === "tool_use") {
-                        openToolAggregateKeysByCallId.set(callId, [...pendingKeys, aggregateKey]);
-                    }
-                    // Inject §N§ prefix into this occurrence's visible text
-                    // when it's a tool_result. (OpenCode parity: prefix
-                    // only goes on the result, not the invocation.)
-                    if (!skipPrefixInjection && part.kind === "tool_result") {
-                        part.setText(prependTag(tagId, text));
-                    }
-                    targets.set(tagId, buildAggregateTarget(tagId, aggregate.occurrences));
                     if (part.kind === "tool_result") {
-                        markToolAggregateResolved(
-                            callId,
-                            aggregateKey,
-                            openToolAggregateKeysByCallId,
-                        );
-                        activeToolResultRun = { callId, aggregateKey };
+                        applyGrownToolResultAccounting({
+                            db,
+                            sessionId,
+                            tagger,
+                            aggregate,
+                            byteSize: accounting.byteSize,
+                            part,
+                            text,
+                            timing,
+                            knownTokenCount: accounting.tokenCount,
+                        });
                     }
+                    syncToolAggregateAccounting(tagger, sessionId, aggregate);
+                    if (timing) timing.identity += performance.now() - identityStart;
+                }
+
+                toolAggregates.set(aggregateKey, aggregate);
+                if (part.kind === "tool_use") {
+                    openToolAggregateKeysByCallId.set(callId, [...pendingKeys, aggregateKey]);
+                }
+                applyToolPrefixAndTarget({
+                    skipPrefixInjection,
+                    part,
+                    text,
+                    tagId: aggregate.tagId,
+                    aggregate,
+                    targets,
+                    timing,
+                });
+                if (part.kind === "tool_result") {
+                    markToolAggregateResolved(callId, aggregateKey, openToolAggregateKeysByCallId);
+                    activeToolResultRun = { callId, aggregateKey };
                 }
             }
             // thinking, image, file, structural, unknown → skip.
         }
     }
 
+    if (timing && options.onTiming) {
+        options.onTiming("identity", timing.identity);
+        options.onTiming("prefix", timing.prefix);
+        options.onTiming("targets", timing.targets);
+        options.onTiming("tokenCounting", timing.tokenCounting);
+    }
     return { targets };
 }
 
+interface AggregateToolAccounting {
+    text: string;
+    byteSize: number;
+    tokenCount: number;
+    toolName: string | null;
+    inputByteSize: number;
+    inputTokenCount: number;
+}
+
+interface GrownToolResultAccountingArgs {
+    db: ContextDatabase;
+    sessionId: string;
+    tagger: Tagger;
+    aggregate: ToolAggregate & { tagId: number };
+    byteSize: number;
+    part: TranscriptPart;
+    text: string;
+    timing: TagTranscriptTiming | undefined;
+    tokenCache?: Map<string, { text: string; tokenCount: number }>;
+    tokenCacheKey?: string;
+    knownTokenCount?: number;
+}
+
+function syncToolAggregateAccounting(
+    tagger: Tagger,
+    sessionId: string,
+    aggregate: ToolAggregate & { tagId: number },
+): void {
+    tagger.setToolTagAccounting(sessionId, aggregate.tagId, {
+        byteSize: aggregate.maxByteSize,
+        tokenCount: aggregate.maxTokenCount,
+        inputByteSize: aggregate.inputByteSize,
+        inputTokenCount: aggregate.inputTokenCount,
+    });
+}
+
+function applyGrownToolResultAccounting(args: GrownToolResultAccountingArgs): void {
+    if (args.byteSize <= args.aggregate.maxByteSize) return;
+
+    let tokenCount = args.knownTokenCount;
+    if (tokenCount !== undefined && args.tokenCacheKey) {
+        args.tokenCache?.set(args.tokenCacheKey, { text: args.text, tokenCount });
+    }
+    if (tokenCount === undefined) {
+        const tokenStart = args.timing ? performance.now() : 0;
+        tokenCount = getCachedToolPartTokenCount(
+            args.part,
+            args.text,
+            args.tokenCache,
+            args.tokenCacheKey,
+        );
+        if (args.timing) args.timing.tokenCounting += performance.now() - tokenStart;
+    }
+    args.aggregate.maxByteSize = args.byteSize;
+    args.aggregate.maxTokenCount = tokenCount;
+    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);
+}
+
+function readAggregateToolAccounting(
+    part: TranscriptPart,
+    timing: TagTranscriptTiming | undefined,
+    tokenCache?: Map<string, { text: string; tokenCount: number }>,
+    tokenCacheKey?: string,
+): AggregateToolAccounting {
+    const text = part.getText() ?? "";
+    const byteSize = getToolPartByteSize(part, text);
+    const metadata = part.getToolMetadata();
+    let tokenCount = 0;
+    if (part.kind === "tool_result") {
+        const tokenStart = timing ? performance.now() : 0;
+        tokenCount = getCachedToolPartTokenCount(part, text, tokenCache, tokenCacheKey);
+        if (timing) timing.tokenCounting += performance.now() - tokenStart;
+    }
+    return {
+        text,
+        byteSize,
+        tokenCount,
+        toolName: metadata.toolName ?? null,
+        inputByteSize: metadata.inputByteSize,
+        inputTokenCount: metadata.inputTokenCount,
+    };
+}
+
+interface ApplyToolPrefixAndTargetArgs {
+    skipPrefixInjection: boolean;
+    part: TranscriptPart;
+    text: string;
+    tagId: number;
+    aggregate: ToolAggregate;
+    targets: Map<number, TagTarget>;
+    timing: TagTranscriptTiming | undefined;
+}
+
+function applyToolPrefixAndTarget(args: ApplyToolPrefixAndTargetArgs): void {
+    if (!args.skipPrefixInjection && args.part.kind === "tool_result") {
+        const prefixStart = args.timing ? performance.now() : 0;
+        args.part.setText(prependTag(args.tagId, args.text));
+        if (args.timing) args.timing.prefix += performance.now() - prefixStart;
+    }
+    const targetStart = args.timing ? performance.now() : 0;
+    args.targets.set(args.tagId, buildAggregateTarget(args.tagId, args.aggregate.occurrences));
+    if (args.timing) args.timing.targets += performance.now() - targetStart;
+}
+
 function findLastUnresolvedToolAggregateKey(
     pendingKeys: string[],
     toolAggregates: Map<string, ToolAggregate & { tagId: number }>,
@@ -466,6 +702,19 @@ function getToolPartTokenCount(part: TranscriptPart, text: string): number {
     return 0;
 }
 
+function getCachedToolPartTokenCount(
+    part: TranscriptPart,
+    text: string,
+    cache?: Map<string, { text: string; tokenCount: number }>,
+    cacheKey?: string,
+): number {
+    const cached = cacheKey ? cache?.get(cacheKey) : undefined;
+    if (cached?.text === text) return cached.tokenCount;
+    const tokenCount = getToolPartTokenCount(part, text);
+    if (cacheKey) cache?.set(cacheKey, { text, tokenCount });
+    return tokenCount;
+}
+
 function getNonTextToolResultByteSize(part: TranscriptPart): number {
     // Prefer the adapter's exact raw-payload size when available (Pi's
     // tool_result proxy can serialize the real content array, incl. images).
@@ -508,11 +757,23 @@ interface TagTextPartArgs {
     targets: Map<number, TagTarget>;
     skipPrefixInjection: boolean;
     entryFingerprint: string | null;
+    reuseIdentity: boolean;
+    timing?: TagTranscriptTiming;
+    textTokenCache?: Map<string, { text: string; tokenCount: number }>;
 }
 
 function tagTextPart(args: TagTextPartArgs): void {
+    const identityStart = args.timing ? performance.now() : 0;
     const text = args.part.getText() ?? "";
     const contentId = `${args.messageId}:p${args.textOrdinal}`;
+    const reusableTagId = args.reuseIdentity
+        ? args.tagger.getTag(args.sessionId, contentId, "message")
+        : undefined;
+    if (reusableTagId !== undefined) {
+        if (args.timing) args.timing.identity += performance.now() - identityStart;
+        applyTextPrefixAndTarget(args, reusableTagId, text);
+        return;
+    }
     const tagId = args.tagger.assignTag(
         args.sessionId,
         contentId,
@@ -525,11 +786,24 @@ function tagTextPart(args: TagTextPartArgs): void {
         args.entryFingerprint,
         // Lazy: fires only on fresh insert. Strip any §N§ prefix so a re-tag
         // from already-prefixed text still tokenizes the pristine content.
-        () => ({
-            tokenCount: estimateTagTextTokens(stripTagPrefix(text)),
-            inputTokenCount: null,
-            reasoningTokenCount: null,
-        }),
+        () => {
+            const tokenStart = args.timing ? performance.now() : 0;
+            const cached = args.textTokenCache?.get(contentId);
+            const tokenCount =
+                cached?.text === text
+                    ? cached.tokenCount
+                    : estimateTagTextTokens(stripTagPrefix(text));
+            if (cached?.text !== text) {
+                args.textTokenCache?.set(contentId, { text, tokenCount });
+            }
+            const counts = {
+                tokenCount,
+                inputTokenCount: null,
+                reasoningTokenCount: null,
+            };
+            if (args.timing) args.timing.tokenCounting += performance.now() - tokenStart;
+            return counts;
+        },
     );
 
     // Persist the original (pre-tagged) source content so caveman
@@ -547,12 +821,20 @@ function tagTextPart(args: TagTextPartArgs): void {
     if (sourceContent.trim().length > 0) {
         saveSourceContent(args.db, args.sessionId, tagId, sourceContent);
     }
+    if (args.timing) args.timing.identity += performance.now() - identityStart;
+    applyTextPrefixAndTarget(args, tagId, text);
+}
 
+function applyTextPrefixAndTarget(args: TagTextPartArgs, tagId: number, text: string): void {
     if (!args.skipPrefixInjection) {
+        const prefixStart = args.timing ? performance.now() : 0;
         args.part.setText(prependTag(tagId, text));
+        if (args.timing) args.timing.prefix += performance.now() - prefixStart;
     }
 
+    const targetStart = args.timing ? performance.now() : 0;
     args.targets.set(tagId, buildTextTarget(args.part, args.message));
+    if (args.timing) args.timing.targets += performance.now() - targetStart;
 }
 
 interface TagToolPartArgs {
@@ -566,9 +848,12 @@ interface TagToolPartArgs {
     db: ContextDatabase;
     targets: Map<number, TagTarget>;
     skipPrefixInjection: boolean;
+    reuseIdentity: boolean;
+    timing?: TagTranscriptTiming;
 }
 
 function tagToolPart(args: TagToolPartArgs): void {
+    const identityStart = args.timing ? performance.now() : 0;
     // Prefer the part's stable id (tool call id from Pi/OpenCode); fall
     // back to a synthetic locator. Tool calls and their results MAY
     // share an id (Pi sets toolCallId on ToolResultMessage to match the
@@ -577,10 +862,21 @@ function tagToolPart(args: TagToolPartArgs): void {
     // target the call-id pair as a unit.
     const stableId = args.part.id;
     const contentId = stableId ?? `${args.messageId}:t${args.partIndex}`;
+    const reusableTagId = args.reuseIdentity
+        ? args.tagger.getToolTag(args.sessionId, contentId, contentId)
+        : undefined;
+    if (reusableTagId !== undefined) {
+        const text = args.part.kind === "tool_result" ? (args.part.getText() ?? "") : "";
+        if (args.timing) args.timing.identity += performance.now() - identityStart;
+        applySingleToolPrefixAndTarget(args, reusableTagId, text);
+        return;
+    }
     const text = args.part.getText() ?? "";
     const toolByteSize = getToolPartByteSize(args.part, text);
-    const toolTokenCount = getToolPartTokenCount(args.part, text);
     const meta = args.part.getToolMetadata();
+    const tokenStart = args.timing ? performance.now() : 0;
+    const toolTokenCount = getToolPartTokenCount(args.part, text);
+    if (args.timing) args.timing.tokenCounting += performance.now() - tokenStart;
     // v3.3.1 Layer C: synthetic ownership for the no-callId Pi
     // fallback. Owner == callId == contentId. The composite key
     // collapses to a unique synthetic identifier per part, preserving
@@ -596,22 +892,34 @@ function tagToolPart(args: TagToolPartArgs): void {
         0,
         meta.toolName ?? null,
         meta.inputByteSize,
-        () => ({
-            tokenCount: toolTokenCount,
-            inputTokenCount: meta.inputTokenCount,
-            reasoningTokenCount: null,
-        }),
+        () => {
+            const tokenStart = args.timing ? performance.now() : 0;
+            const counts = {
+                tokenCount: toolTokenCount,
+                inputTokenCount: meta.inputTokenCount,
+                reasoningTokenCount: null,
+            };
+            if (args.timing) args.timing.tokenCounting += performance.now() - tokenStart;
+            return counts;
+        },
     );
+    if (args.timing) args.timing.identity += performance.now() - identityStart;
+    applySingleToolPrefixAndTarget(args, tagId, text);
+}
 
+function applySingleToolPrefixAndTarget(args: TagToolPartArgs, tagId: number, text: string): void {
     // For tool parts, the visible payload is the tool result text. We
     // can inject the tag prefix into it for in-text references; this
     // matches the OpenCode behavior of tagging tool outputs.
     if (!args.skipPrefixInjection && args.part.kind === "tool_result") {
-        const tagged = prependTag(tagId, text);
-        args.part.setText(tagged);
+        const prefixStart = args.timing ? performance.now() : 0;
+        args.part.setText(prependTag(tagId, text));
+        if (args.timing) args.timing.prefix += performance.now() - prefixStart;
     }
 
+    const targetStart = args.timing ? performance.now() : 0;
     args.targets.set(tagId, buildToolTarget(args.part, args.message, tagId));
+    if (args.timing) args.timing.targets += performance.now() - targetStart;
 }
 
 function setToolContentOrText(part: TranscriptPart, content: string): boolean {