@cortexkit/opencode-magic-context 0.21.7 → 0.22.0

package/src/shared/tag-transcript.ts

@@ -48,7 +48,7 @@
 import type { ContextDatabase } from "../features/magic-context/storage";
 import { saveSourceContent } from "../features/magic-context/storage-source";
 import { updateTagByteSize, updateTagInputByteSize } from "../features/magic-context/storage-tags";
-import type { Tagger } from "../features/magic-context/tagger";
+import { makeToolCompositeKey, type Tagger } from "../features/magic-context/tagger";
 import {
     byteSize,
     prependTag,
@@ -67,6 +67,15 @@ export interface TagTranscriptOptions {
      * consistent across passes.
      */
     skipPrefixInjection?: boolean;
+    /**
+     * Pi-only: map of messageId → raw-message fingerprint. When a NEW message
+     * text tag is created, its fingerprint is persisted on the tag row so a
+     * later pass can adopt the fallback-id tag onto the real SessionEntry id
+     * (keeping tag_number/§N§ stable). OpenCode omits this → tags store NULL
+     * → adoption never fires. Keyed by the bare messageId (not the `:pN`
+     * contentId) since all parts of a message share one fingerprint.
+     */
+    entryFingerprintByMessageId?: ReadonlyMap<string, string>;
 }
 
 export interface TagTranscriptResult {
@@ -138,13 +147,13 @@ export function tagTranscript(
     const skipPrefixInjection = options.skipPrefixInjection === true;
     const targets = new Map<number, TagTarget>();
 
-    // Per-callId tool aggregation tracked across the single walk. Tags
-    // get allocated in transcript order (first occurrence reserves the
-    // tag number); subsequent occurrences reuse the same tag and merge
-    // their occurrence into the aggregate's TagTarget. This preserves
-    // chronological tag numbering while still aggregating drop behavior
-    // across both invocation and result.
+    // Tool aggregation is keyed by the same owner+callId identity used by
+    // assignToolTag. OpenCode/Pi callId counters can repeat across turns, so
+    // a bare callId key can merge distinct invocations and replay drops/status
+    // changes against the wrong tool pair.
     const toolAggregates = new Map<string, ToolAggregate & { tagId: number }>();
+    const openToolAggregateKeysByCallId = new Map<string, string[]>();
+    let activeToolResultRun: { callId: string; aggregateKey: string } | undefined;
 
     // v3.3.1 Layer C (plan v3.3.1 Finding #16): the previous outer
     // db.transaction() wrapper rolled back EVERY tag insert + savedSource
@@ -156,6 +165,7 @@ export function tagTranscript(
     for (let msgIndex = 0; msgIndex < transcript.messages.length; msgIndex += 1) {
         const message = transcript.messages[msgIndex];
         if (message === undefined) continue;
+        activeToolResultRun = undefined;
         const messageId = message.info.id;
 
         let textOrdinal = 0;
@@ -165,6 +175,10 @@ export function tagTranscript(
             const part = parts[partIndex];
             if (part === undefined) continue;
 
+            if (part.kind !== "tool_result") {
+                activeToolResultRun = undefined;
+            }
+
             if (part.kind === "text") {
                 // Synthetic message ids (Pi tail synthetic user with
                 // no id) cannot be tagged — there's no stable handle
@@ -186,19 +200,25 @@ export function tagTranscript(
                     db,
                     targets,
                     skipPrefixInjection,
+                    entryFingerprint: options.entryFingerprintByMessageId?.get(messageId) ?? null,
                 });
                 textOrdinal += 1;
                 continue;
             }
 
             if (part.kind === "tool_use" || part.kind === "tool_result") {
-                if (messageId === undefined) continue;
+                if (messageId === undefined) {
+                    activeToolResultRun = undefined;
+                    continue;
+                }
 
                 const callId = part.id;
                 const text = part.getText() ?? "";
+                const toolByteSize = getToolPartByteSize(part, text);
                 const meta = part.getToolMetadata();
 
                 if (typeof callId !== "string" || callId.length === 0) {
+                    activeToolResultRun = undefined;
                     // No stable callId to aggregate on. Tag independently.
                     tagToolPart({
                         sessionId,
@@ -215,21 +235,35 @@ export function tagTranscript(
                     continue;
                 }
 
-                const existing = toolAggregates.get(callId);
+                const pendingKeys = openToolAggregateKeysByCallId.get(callId) ?? [];
+                let existingKey: string | undefined;
+                if (part.kind === "tool_result") {
+                    if (
+                        activeToolResultRun !== undefined &&
+                        activeToolResultRun.callId === callId
+                    ) {
+                        existingKey = activeToolResultRun.aggregateKey;
+                    } else {
+                        existingKey = findLastUnresolvedToolAggregateKey(
+                            pendingKeys,
+                            toolAggregates,
+                        );
+                    }
+                }
+                const aggregateKey: string = existingKey ?? makeToolCompositeKey(messageId, callId);
+                const existing = toolAggregates.get(aggregateKey);
                 if (existing) {
-                    // Second (or later) occurrence for this call_id.
-                    // Merge into the existing aggregate, update byte_size
-                    // in DB if larger, and rebuild the TagTarget so the
-                    // closures over `occurrences` see all parts.
+                    // 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,
                         part,
                         kind: part.kind,
                     });
-                    const newByteSize = byteSize(text);
-                    if (newByteSize > existing.maxByteSize) {
-                        existing.maxByteSize = newByteSize;
-                        updateTagByteSize(db, sessionId, existing.tagId, newByteSize);
+                    if (toolByteSize > existing.maxByteSize) {
+                        existing.maxByteSize = toolByteSize;
+                        updateTagByteSize(db, sessionId, existing.tagId, toolByteSize);
                     }
                     if (existing.toolName === null && meta.toolName) {
                         existing.toolName = meta.toolName;
@@ -253,18 +287,23 @@ export function tagTranscript(
                         existing.tagId,
                         buildAggregateTarget(existing.tagId, existing.occurrences),
                     );
+                    if (part.kind === "tool_result") {
+                        markToolAggregateResolved(
+                            callId,
+                            aggregateKey,
+                            openToolAggregateKeysByCallId,
+                        );
+                        activeToolResultRun = { callId, aggregateKey };
+                    }
                 } else {
-                    // First occurrence — reserve the tag number.
-                    // v3.3.1 Layer C: Pi main aggregation path. Owner
-                    // is the Pi message hosting the tool aggregate.
-                    // Owner stays stable across passes because Pi
-                    // re-emits the full transcript each time and
-                    // message ids are durable.
+                    // First occurrence for this owner+callId identity — reserve
+                    // the tag number. Owner stays stable across passes because
+                    // transcript message ids are durable.
                     const tagId = tagger.assignToolTag(
                         sessionId,
                         callId,
                         messageId,
-                        byteSize(text),
+                        toolByteSize,
                         db,
                         0,
                         meta.toolName ?? null,
@@ -280,11 +319,14 @@ export function tagTranscript(
                                 kind: part.kind,
                             },
                         ],
-                        maxByteSize: byteSize(text),
+                        maxByteSize: toolByteSize,
                         toolName: meta.toolName ?? null,
                         inputByteSize: part.kind === "tool_use" ? meta.inputByteSize : 0,
                     };
-                    toolAggregates.set(callId, aggregate);
+                    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.)
@@ -292,6 +334,14 @@ export function tagTranscript(
                         part.setText(prependTag(tagId, text));
                     }
                     targets.set(tagId, buildAggregateTarget(tagId, aggregate.occurrences));
+                    if (part.kind === "tool_result") {
+                        markToolAggregateResolved(
+                            callId,
+                            aggregateKey,
+                            openToolAggregateKeysByCallId,
+                        );
+                        activeToolResultRun = { callId, aggregateKey };
+                    }
                 }
             }
             // thinking, image, file, structural, unknown → skip.
@@ -301,6 +351,69 @@ export function tagTranscript(
     return { targets };
 }
 
+function findLastUnresolvedToolAggregateKey(
+    pendingKeys: string[],
+    toolAggregates: Map<string, ToolAggregate & { tagId: number }>,
+): string | undefined {
+    for (let i = pendingKeys.length - 1; i >= 0; i -= 1) {
+        const key = pendingKeys[i];
+        if (key === undefined) continue;
+        const aggregate = toolAggregates.get(key);
+        if (aggregate === undefined) continue;
+        if (!aggregate.occurrences.some((occ) => occ.kind === "tool_result")) {
+            return key;
+        }
+    }
+    return undefined;
+}
+
+function markToolAggregateResolved(
+    callId: string,
+    aggregateKey: string,
+    openToolAggregateKeysByCallId: Map<string, string[]>,
+): void {
+    const pendingKeys = openToolAggregateKeysByCallId.get(callId);
+    if (pendingKeys === undefined) return;
+    const nextPendingKeys = pendingKeys.filter((key) => key !== aggregateKey);
+    if (nextPendingKeys.length === 0) {
+        openToolAggregateKeysByCallId.delete(callId);
+        return;
+    }
+    openToolAggregateKeysByCallId.set(callId, nextPendingKeys);
+}
+
+function getToolPartByteSize(part: TranscriptPart, text: string): number {
+    const textByteSize = byteSize(text);
+    if (textByteSize > 0 || part.kind !== "tool_result") return textByteSize;
+    return getNonTextToolResultByteSize(part);
+}
+
+function getNonTextToolResultByteSize(part: TranscriptPart): number {
+    const record = isRecord(part) ? part : undefined;
+    const content =
+        record?.content ??
+        record?.rawContent ??
+        record?.rawPart ??
+        record?.part ??
+        record?.data ??
+        record?.image ??
+        record?.source;
+    const serialized = safeJsonStringify(content ?? part);
+    return serialized === undefined ? 0 : byteSize(serialized);
+}
+
+function safeJsonStringify(value: unknown): string | undefined {
+    try {
+        return JSON.stringify(value);
+    } catch {
+        return undefined;
+    }
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+    return typeof value === "object" && value !== null;
+}
+
 interface TagTextPartArgs {
     sessionId: string;
     message: { info: { id?: string; role: string } };
@@ -312,6 +425,7 @@ interface TagTextPartArgs {
     db: ContextDatabase;
     targets: Map<number, TagTarget>;
     skipPrefixInjection: boolean;
+    entryFingerprint: string | null;
 }
 
 function tagTextPart(args: TagTextPartArgs): void {
@@ -323,6 +437,10 @@ function tagTextPart(args: TagTextPartArgs): void {
         "message",
         byteSize(text),
         args.db,
+        0,
+        null,
+        0,
+        args.entryFingerprint,
     );
 
     // Persist the original (pre-tagged) source content so caveman
@@ -371,6 +489,7 @@ function tagToolPart(args: TagToolPartArgs): void {
     const stableId = args.part.id;
     const contentId = stableId ?? `${args.messageId}:t${args.partIndex}`;
     const text = args.part.getText() ?? "";
+    const toolByteSize = getToolPartByteSize(args.part, text);
     const meta = args.part.getToolMetadata();
     // v3.3.1 Layer C: synthetic ownership for the no-callId Pi
     // fallback. Owner == callId == contentId. The composite key
@@ -382,7 +501,7 @@ function tagToolPart(args: TagToolPartArgs): void {
         args.sessionId,
         contentId,
         contentId,
-        byteSize(text),
+        toolByteSize,
         args.db,
         0,
         meta.toolName ?? null,
@@ -400,6 +519,18 @@ function tagToolPart(args: TagToolPartArgs): void {
     args.targets.set(tagId, buildToolTarget(args.part, args.message));
 }
 
+function setToolContentOrText(part: TranscriptPart, content: string): boolean {
+    try {
+        if (part.setToolOutput(content)) return true;
+    } catch {
+        // Pi assistant tool_use parts deliberately assert if callers try
+        // to write a nonexistent output slot. Truncated-mode drops still
+        // need to shrink the invocation, so fall back to visible text/args
+        // replacement while preserving the adapter-level invariant.
+    }
+    return part.setText(content);
+}
+
 /**
  * Build a TagTarget that walks ALL occurrences of a tool call (invocation
  * + result) when mutating. This is the per-callId aggregate target used
@@ -428,9 +559,7 @@ function buildAggregateTarget(tagId: number, occurrences: ToolOccurrence[]): Tag
             for (const occ of occurrences) {
                 // Try setToolOutput first (works on tool_result-shaped parts);
                 // fall back to setText so tool_use parts also get sentinelized.
-                if (occ.part.setToolOutput(content)) {
-                    changed = true;
-                } else if (occ.part.setText(content)) {
+                if (setToolContentOrText(occ.part, content)) {
                     changed = true;
                 }
             }
@@ -461,7 +590,7 @@ function buildAggregateTarget(tagId: number, occurrences: ToolOccurrence[]): Tag
             const sentinel = "[truncated]";
             let any = false;
             for (const occ of occurrences) {
-                if (occ.part.setToolOutput(sentinel) || occ.part.setText(sentinel)) {
+                if (setToolContentOrText(occ.part, sentinel)) {
                     any = true;
                 }
             }
@@ -521,7 +650,7 @@ function buildToolTarget(
 ): TagTarget {
     return {
         setContent(content: string): boolean {
-            return part.setToolOutput(content) || part.setText(content);
+            return setToolContentOrText(part, content);
         },
         getContent(): string | null {
             return part.getText() ?? null;
@@ -542,7 +671,7 @@ function buildToolTarget(
             // via setToolOutput so the underlying tool_result content
             // gets the truncation; falls back to setText for cases
             // where the part type doesn't support setToolOutput.
-            const ok = part.setToolOutput("[truncated]") || part.setText("[truncated]");
+            const ok = setToolContentOrText(part, "[truncated]");
             return ok ? "truncated" : "absent";
         },
         message: {

package/src/tui/data/context-db.ts

@@ -10,7 +10,8 @@ import type { RpcNotificationMessage, SidebarSnapshot, StatusDetail } from "../.
 export type { SidebarSnapshot, StatusDetail };
 
 let rpcClient: MagicContextRpcClient | null = null;
-let lastReceivedNotificationId = 0;
+let rpcGeneration = 0;
+const lastReceivedNotificationIdBySession = new Map<string, number>();
 
 function getStorageDir(): string {
     // Plugin v0.16+ uses the shared cortexkit/magic-context path so OpenCode
@@ -24,14 +25,25 @@ function getStorageDir(): string {
 /** Initialize the RPC client. Call once on TUI startup. */
 export function initRpcClient(directory: string): void {
     const storageDir = getStorageDir();
+    // Bump the generation before replacing the client so late notification
+    // responses from a disposed client cannot repopulate cleared cursors.
+    rpcGeneration += 1;
+    lastReceivedNotificationIdBySession.clear();
     rpcClient = new MagicContextRpcClient(storageDir, directory);
 }
 
+export function getRpcGeneration(): number {
+    return rpcGeneration;
+}
+
 /** Clean up the RPC client. */
 export function closeRpc(): void {
+    // Closing invalidates any already-issued RPC calls; their callbacks must
+    // observe the new generation and avoid advancing stale notification cursors.
+    rpcGeneration += 1;
     rpcClient?.reset();
     rpcClient = null;
-    lastReceivedNotificationId = 0;
+    lastReceivedNotificationIdBySession.clear();
 }
 
 const EMPTY_SNAPSHOT: SidebarSnapshot = {
@@ -55,10 +67,14 @@ const EMPTY_SNAPSHOT: SidebarSnapshot = {
     compartmentTokens: 0,
     factTokens: 0,
     memoryTokens: 0,
+    docsTokens: 0,
+    profileTokens: 0,
     conversationTokens: 0,
     toolCallTokens: 0,
     toolDefinitionTokens: 0,
     executeThreshold: 65,
+    newWorkTokens: null,
+    totalInputTokens: null,
 };
 
 /**
@@ -227,7 +243,35 @@ export async function requestRecomp(sessionId: string): Promise<boolean> {
     }
 }
 
+/** Run `/ctx-session-upgrade` for the session (full recomp + once-per-project
+ *  memory migration). Fired from the upgrade dialog's "Run upgrade now" action. */
+export async function requestUpgrade(sessionId: string): Promise<boolean> {
+    if (!rpcClient) return false;
+    try {
+        const result = await rpcClient.call<{ ok: boolean }>("upgrade", { sessionId });
+        return result.ok ?? false;
+    } catch {
+        return false;
+    }
+}
+
+/** Mark the upgrade reminder dismissed (the user made an explicit Confirm/Cancel
+ *  choice), setting the durable stamp so the FRESH dialog won't re-show. Resume
+ *  prompts are staging-driven and unaffected. */
+export async function dismissUpgradeReminder(sessionId: string): Promise<boolean> {
+    if (!rpcClient) return false;
+    try {
+        const result = await rpcClient.call<{ ok: boolean }>("dismiss-upgrade-reminder", {
+            sessionId,
+        });
+        return result.ok ?? false;
+    } catch {
+        return false;
+    }
+}
+
 export interface TuiMessage {
+    id: number;
     type: string;
     payload: Record<string, unknown>;
     sessionId?: string;
@@ -277,20 +321,24 @@ export async function markAnnounced(): Promise<boolean> {
 }
 
 /** Poll for pending server→TUI notifications via RPC. */
-export async function consumeTuiMessages(): Promise<TuiMessage[]> {
+export async function consumeTuiMessages(sessionId: string): Promise<TuiMessage[]> {
     if (!rpcClient) return [];
     try {
         const result = await rpcClient.call<{ messages: RpcNotificationMessage[] }>(
             "pending-notifications",
-            { lastReceivedId: lastReceivedNotificationId },
+            // Pass the TUI's active session so the server only drains
+            // notifications scoped to it (or global ones). Without this, a
+            // notification for another session served by the same process (e.g.
+            // OpenCode Desktop on the same project) could surface here. The
+            // cursor is per-session and is advanced by the poller only after it
+            // has delivered the returned batch.
+            {
+                lastReceivedId: lastReceivedNotificationIdBySession.get(sessionId) ?? 0,
+                sessionId,
+            },
         );
-        const messages = result.messages ?? [];
-        for (const message of messages) {
-            if (message.id > lastReceivedNotificationId) {
-                lastReceivedNotificationId = message.id;
-            }
-        }
-        return messages.map((m) => ({
+        return (result.messages ?? []).map((m) => ({
+            id: m.id,
             type: m.type,
             payload: m.payload,
             sessionId: m.sessionId,
@@ -299,3 +347,21 @@ export async function consumeTuiMessages(): Promise<TuiMessage[]> {
         return [];
     }
 }
+
+/**
+ * Advance the delivered-message cursor for one active TUI session.
+ * Callers must pass only the contiguous handled prefix of the drained batch;
+ * this helper remains empty-safe and monotonic for that prefix.
+ */
+export function markTuiMessagesHandled(sessionId: string, messages: TuiMessage[]): void {
+    const previous = lastReceivedNotificationIdBySession.get(sessionId) ?? 0;
+    let next = previous;
+    for (const message of messages) {
+        if (message.id > next) {
+            next = message.id;
+        }
+    }
+    if (next > previous) {
+        lastReceivedNotificationIdBySession.set(sessionId, next);
+    }
+}