From 77f8dda4a902693a7d5df1db33287469ca55968b Mon Sep 17 00:00:00 2001 From: SilentBless <126187696+SilentBless@users.noreply.github.com> Date: Wed, 9 Sep 2026 21:13:30 +0300 Subject: [PATCH 1/7] fix(pi): reduce OMP native inputs and optional Codex reasoning --- packages/pi-plugin/README.md | 2 + .../pi-plugin/src/context-handler.test.ts | 37 ++- packages/pi-plugin/src/context-handler.ts | 20 +- .../pi-plugin/src/native-replay-pi.test.ts | 308 ++++++++++++++++++ packages/pi-plugin/src/native-replay-pi.ts | 187 +++++++++++ .../pi-plugin/src/reasoning-replay-pi.test.ts | 113 +++++++ packages/pi-plugin/src/reasoning-replay-pi.ts | 19 ++ .../pi-plugin/src/signal-peek-drain.test.ts | 6 - packages/pi-plugin/src/transcript-pi.test.ts | 85 +++++ packages/pi-plugin/src/transcript-pi.ts | 11 + 10 files changed, 778 insertions(+), 10 deletions(-) create mode 100644 packages/pi-plugin/src/native-replay-pi.test.ts create mode 100644 packages/pi-plugin/src/native-replay-pi.ts diff --git a/packages/pi-plugin/README.md b/packages/pi-plugin/README.md index ed8e9e549..f6a7c35da 100644 --- a/packages/pi-plugin/README.md +++ b/packages/pi-plugin/README.md @@ -76,6 +76,8 @@ npx @cortexkit/magic-context@latest doctor --harness omp OMP's legacy Pi loader maps `@earendil-works/*` imports to its bundled `@oh-my-pi/*` runtime. Magic Context identifies OMP from `@oh-my-pi/pi-utils`'s in-process `APP_NAME` through the running host's module graph, not from an executable basename, so compiled and symlinked launches behave the same. It also re-invokes the current host executable for historian, dreamer, and sidekick children, so OMP children remain OMP processes. +OMP can replay Responses/Codex history from `providerPayload` instead of ordinary message content. Tool-input reductions update the matched native call without discarding unrelated native results. Old encrypted reasoning is cleared only for Codex models whose resolved compatibility settings explicitly allow omission, and only from incremental history. Other models, full snapshots, plaintext or malformed reasoning, redacted thinking, and computer-linked reasoning retain their native reasoning. Native text and user/developer history carriers are not rewritten. + --- ## Configuration diff --git a/packages/pi-plugin/src/context-handler.test.ts b/packages/pi-plugin/src/context-handler.test.ts index 603c52c8a..59c2a630a 100644 --- a/packages/pi-plugin/src/context-handler.test.ts +++ b/packages/pi-plugin/src/context-handler.test.ts @@ -3622,8 +3622,10 @@ describe("registerPiContextHandler", () => { it("restores reasoning bytes when the durable watermark write fails", async () => { const db = createTestDb(); const sessionId = "ses-reasoning-watermark-failure"; + let watermarkWriteAttempted = false; const restorePersistence = contextHandlerInternals.setReasoningWatermarkPersistenceForTests(() => { + watermarkWriteAttempted = true; throw new Error("faulted reasoning watermark write"); }); try { @@ -3642,6 +3644,16 @@ describe("registerPiContextHandler", () => { { role: "assistant", timestamp: 2, + providerPayload: { + type: "openaiResponsesHistory", + dt: true, + items: [ + { + type: "reasoning", + encrypted_content: "durable native reasoning", + }, + ], + }, content: [ { type: "thinking", @@ -3664,10 +3676,20 @@ describe("registerPiContextHandler", () => { messages as never, ), getContextUsage: () => ({ - tokens: 70_000, - percent: 70, + tokens: 90_000, + percent: 90, contextWindow: 100_000, }), + model: { + id: "test-codex", + api: "openai-codex-responses", + provider: "openai-codex", + contextWindow: 100_000, + compat: { + requiresReasoningContentForAllAssistantTurns: false, + requiresReasoningContentForToolCalls: false, + }, + }, } as never); if (!result) throw new Error("expected transformed messages"); return result.messages; @@ -3675,12 +3697,23 @@ describe("registerPiContextHandler", () => { const first = await runPass(); const second = await runPass(); + expect(watermarkWriteAttempted).toBe(true); const firstThinking = (first[1] as { content: Record[] }) .content[0]; expect(firstThinking).toMatchObject({ thinking: "durable secret", thinkingSignature: "sig", }); + expect(first[1]).toMatchObject({ + providerPayload: { + items: [ + { + type: "reasoning", + encrypted_content: "durable native reasoning", + }, + ], + }, + }); expect(JSON.stringify(second)).toBe(JSON.stringify(first)); expect( getOrCreateSessionMeta(db, sessionId).clearedReasoningThroughTag, diff --git a/packages/pi-plugin/src/context-handler.ts b/packages/pi-plugin/src/context-handler.ts index f5421abab..27088f6d2 100644 --- a/packages/pi-plugin/src/context-handler.ts +++ b/packages/pi-plugin/src/context-handler.ts @@ -223,6 +223,7 @@ import { prepareCachedM0M1PiReplay, trimPiMessagesToCachedBoundary, } from "./inject-compartments-pi"; +import { canClearNativeReasoning } from "./native-replay-pi"; import { hasVisibleNoteReadCallPi } from "./note-visibility-pi"; import { resolvePiUsableContextLimit, @@ -3086,6 +3087,7 @@ export function registerPiContextHandler( clearReasoningAge: options.heuristics?.clearReasoningAge ?? DEFAULT_CLEAR_REASONING_AGE, + nativeReasoningMayClear: canClearNativeReasoning(ctx.model), }, canUseEmptySentinels, temporalAwareness: options.injection?.temporalAwareness === true, @@ -4529,6 +4531,7 @@ interface RunPipelineArgs { */ reasoningClearing?: { clearReasoningAge: number; + nativeReasoningMayClear: boolean; }; /** True only when the active provider filters empty sentinel content safely. */ canUseEmptySentinels: boolean; @@ -4624,16 +4627,27 @@ function captureReasoningMutationRollback( ): () => void { const snapshots: Array<{ part: Record; - field: "thinking" | "text"; + field: "thinking" | "text" | "providerPayload"; value: unknown; hadSignature?: boolean; signature?: unknown; }> = []; for (const raw of messages) { if (!raw || typeof raw !== "object") continue; - const message = raw as { role?: unknown; content?: unknown }; + const message = raw as { + role?: unknown; + content?: unknown; + providerPayload?: unknown; + }; if (message.role !== "assistant" || !Array.isArray(message.content)) continue; + if (Object.hasOwn(message, "providerPayload")) { + snapshots.push({ + part: message, + field: "providerPayload", + value: message.providerPayload, + }); + } for (const rawPart of message.content) { if (!rawPart || typeof rawPart !== "object") continue; const part = rawPart as Record; @@ -5339,6 +5353,7 @@ async function runPipeline(args: RunPipelineArgs): Promise { sessionId: args.sessionId, messages: workingMessages, messageIdToMaxTag, + nativeReasoningMayClear: args.reasoningClearing.nativeReasoningMayClear, piMessageStableId: stableIdResolver, }); const inlineReplay = replayStrippedInlineThinkingPi({ @@ -5651,6 +5666,7 @@ async function runPipeline(args: RunPipelineArgs): Promise { messages: workingMessages, messageIdToMaxTag, clearReasoningAge: args.reasoningClearing.clearReasoningAge, + nativeReasoningMayClear: args.reasoningClearing.nativeReasoningMayClear, piMessageStableId: stableIdResolver, }); const stripOutcome = stripInlineThinkingPi({ diff --git a/packages/pi-plugin/src/native-replay-pi.test.ts b/packages/pi-plugin/src/native-replay-pi.test.ts new file mode 100644 index 000000000..a511cb1b0 --- /dev/null +++ b/packages/pi-plugin/src/native-replay-pi.test.ts @@ -0,0 +1,308 @@ +import { describe, expect, it } from "bun:test"; +import { + canClearNativeReasoning, + clearNativeReasoning, + rewriteNativeToolInput, +} from "./native-replay-pi"; + +type NativeItem = Record; + +function nativeMessage(items: NativeItem[], dt = true) { + return { + role: "assistant", + providerPayload: { + type: "openaiResponsesHistory", + provider: "openai-codex", + dt, + items, + }, + }; +} + +const optionalReasoning = { + requiresReasoningContentForAllAssistantTurns: false, + requiresReasoningContentForToolCalls: false, +}; + +describe("native tool input reductions", () => { + it("changes only the selected function input without losing native-only history", () => { + const reasoning = { + type: "reasoning", + encrypted_content: "retain-encrypted-history", + }; + const image = { type: "image_generation_call", result: "retain-image" }; + const text = { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "retain explanation" }], + }; + const call = { + type: "function_call", + id: "fc1", + call_id: "call1", + name: "read", + arguments: JSON.stringify({ path: "original.txt" }), + }; + const message = nativeMessage([reasoning, image, text, call]); + const original = message.providerPayload; + const before = structuredClone(original); + + rewriteNativeToolInput(message, "call1|fc1", { + dropped: "[dropped input]", + }); + + expect(message.providerPayload.items).toEqual([ + reasoning, + image, + text, + { ...call, arguments: JSON.stringify({ dropped: "[dropped input]" }) }, + ]); + expect(original).toEqual(before); + expect(message.providerPayload.provider).toBe(original.provider); + }); + + it("uses the existing OMP custom-input fallback instead of replaying stale raw input", () => { + const call = { + type: "custom_tool_call", + id: "ctc1", + call_id: "call1", + name: "apply_patch", + input: "original patch", + }; + const message = nativeMessage([call]); + rewriteNativeToolInput(message, "call1|ctc1", { + input: "replacement patch", + }); + expect(message.providerPayload.items[0]).toEqual({ + ...call, + input: "replacement patch", + }); + rewriteNativeToolInput(message, "call1|ctc1", { + dropped: "[dropped input]", + }); + expect(message.providerPayload.items[0]).toEqual({ ...call, input: "" }); + expect(call.input).toBe("original patch"); + }); + + it("matches unique id-less calls when OMP supplies a synthesized block id", () => { + const call = { + type: "function_call", + call_id: "function1", + name: "read", + arguments: JSON.stringify({ path: "original.txt" }), + }; + const custom = { + type: "custom_tool_call", + call_id: "custom1", + name: "apply_patch", + input: "original patch", + }; + const message = nativeMessage([call, custom]); + rewriteNativeToolInput(message, "function1|fc_synthesized", { + dropped: "[dropped]", + }); + rewriteNativeToolInput(message, "custom1|fc_synthesized", { + dropped: "[dropped]", + }); + expect(message.providerPayload.items).toEqual([ + { ...call, arguments: JSON.stringify({ dropped: "[dropped]" }) }, + { ...custom, input: "" }, + ]); + expect(message.providerPayload.items[0]).not.toHaveProperty("id"); + expect(message.providerPayload.items[1]).not.toHaveProperty("id"); + }); + + it("does not guess between duplicate, mismatched, or ambiguous id-less calls", () => { + const call = { + type: "function_call", + id: "fc1", + call_id: "call1", + name: "read", + arguments: "{}", + }; + const duplicate = nativeMessage([call, { ...call }]); + const mismatched = nativeMessage([call]); + const ambiguous = nativeMessage([ + call, + { + type: "function_call", + call_id: "call1", + name: "read", + arguments: "{}", + }, + ]); + for (const message of [duplicate, mismatched, ambiguous]) { + const payload = message.providerPayload; + rewriteNativeToolInput(message, "call1|fc_other", { + dropped: "[dropped]", + }); + expect(message.providerPayload).toBe(payload); + } + const payload = duplicate.providerPayload; + rewriteNativeToolInput(duplicate, "call1|fc1", { dropped: "[dropped]" }); + expect(duplicate.providerPayload).toBe(payload); + }); + + it("preserves the replay object when the input is already reduced", () => { + const input = { dropped: "[dropped]" }; + const message = nativeMessage([ + { + type: "function_call", + id: "fc1", + call_id: "call1", + name: "read", + arguments: JSON.stringify(input), + }, + ]); + const payload = message.providerPayload; + rewriteNativeToolInput(message, "call1", input); + expect(message.providerPayload).toBe(payload); + }); + + it("reduces a selected snapshot call without replacing the snapshot or its other history", () => { + const history = { + type: "message", + role: "user", + content: [{ type: "input_text", text: "authoritative history" }], + }; + const call = { + type: "function_call", + id: "fc1", + call_id: "call1", + name: "read", + arguments: "{}", + }; + const message = nativeMessage([history, call], false); + rewriteNativeToolInput(message, "call1|fc1", { dropped: "[dropped]" }); + expect(message.providerPayload).toEqual({ + type: "openaiResponsesHistory", + provider: "openai-codex", + dt: false, + items: [ + history, + { ...call, arguments: JSON.stringify({ dropped: "[dropped]" }) }, + ], + }); + }); +}); + +describe("native reasoning retention", () => { + it("permits clearing only when the Codex model explicitly allows reasoning omission", () => { + expect( + canClearNativeReasoning({ + api: "openai-codex-responses", + compat: optionalReasoning, + }), + ).toBe(true); + expect(canClearNativeReasoning(undefined)).toBe(false); + expect( + canClearNativeReasoning({ api: "openai-codex-responses", compat: {} }), + ).toBe(false); + expect( + canClearNativeReasoning({ + api: "openai-responses", + compat: optionalReasoning, + }), + ).toBe(false); + for (const compat of [ + { + ...optionalReasoning, + requiresReasoningContentForAllAssistantTurns: true, + }, + { ...optionalReasoning, requiresReasoningContentForToolCalls: true }, + { + ...optionalReasoning, + whenThinking: { + ...optionalReasoning, + requiresReasoningContentForAllAssistantTurns: true, + }, + }, + { ...optionalReasoning, whenThinking: {} }, + ]) { + expect( + canClearNativeReasoning({ api: "openai-codex-responses", compat }), + ).toBe(false); + } + }); + + it("removes eligible opaque reasoning without dropping other native items or changing source history", () => { + const reasoning = { + type: "reasoning", + encrypted_content: "opaque-encrypted-content", + content: [], + }; + const text = { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "keep" }], + }; + const call = { + type: "function_call", + call_id: "call1", + name: "read", + arguments: "{}", + }; + const image = { type: "image_generation_call", result: "keep" }; + const message = nativeMessage([reasoning, text, call, image]); + const original = message.providerPayload; + expect(clearNativeReasoning(message, true)).toBe("cleared"); + expect(message.providerPayload.items).toEqual([text, call, image]); + expect(original.items).toEqual([reasoning, text, call, image]); + }); + + it("retains native reasoning when the active model has not authorized omission", () => { + const message = nativeMessage([ + { type: "reasoning", encrypted_content: "must-stay" }, + ]); + const payload = message.providerPayload; + expect(clearNativeReasoning(message, false)).toBe("preserved"); + expect(message.providerPayload).toBe(payload); + }); + + it("keeps snapshots, plaintext, malformed, redacted and computer-linked reasoning intact", () => { + const encrypted = { type: "reasoning", encrypted_content: "keep" }; + const cases = [ + nativeMessage([encrypted], false), + nativeMessage([ + { + ...encrypted, + content: [{ type: "reasoning_text", text: "required plaintext" }], + }, + ]), + nativeMessage([ + encrypted, + { + type: "reasoning", + content: [{ type: "reasoning_text", text: "required" }], + }, + ]), + nativeMessage([{ type: "reasoning", encrypted_content: "" }]), + nativeMessage([ + encrypted, + { type: "computer_call", call_id: "computer1" }, + ]), + { + ...nativeMessage([encrypted]), + content: [{ type: "thinking", thinking: "opaque", redacted: true }], + }, + ]; + for (const message of cases) { + const payload = message.providerPayload; + expect(clearNativeReasoning(message, true)).toBe("preserved"); + expect(message.providerPayload).toBe(payload); + } + }); + + it("leaves ordinary Pi messages and other provider payloads unchanged", () => { + const message = { + role: "assistant", + content: [{ type: "thinking", thinking: "ordinary Pi reasoning" }], + providerPayload: { type: "anthropicMessage", content: "keep" }, + }; + const before = structuredClone(message); + expect(clearNativeReasoning(message, true)).toBe("not-native"); + rewriteNativeToolInput(message, "call1", { dropped: "[dropped]" }); + expect(message).toEqual(before); + expect(clearNativeReasoning(nativeMessage([]), true)).toBe("not-native"); + }); +}); diff --git a/packages/pi-plugin/src/native-replay-pi.ts b/packages/pi-plugin/src/native-replay-pi.ts new file mode 100644 index 000000000..574d8d9ad --- /dev/null +++ b/packages/pi-plugin/src/native-replay-pi.ts @@ -0,0 +1,187 @@ +import { isRecord } from "@magic-context/core/shared/record-type-guard"; + +type NativeEnvelope = { + message: Record; + payload: Record; + items: unknown[]; +}; + +type ToolCallIdentity = { + callId: string; + itemId?: string; +}; + +type ToolCallMatch = { + index: number; + item: Record; + kind: "function_call" | "custom_tool_call"; +}; + +function getNativeEnvelope(message: unknown): NativeEnvelope | undefined { + if (!isRecord(message)) return undefined; + const payload = message.providerPayload; + if ( + !isRecord(payload) || + payload.type !== "openaiResponsesHistory" || + !Array.isArray(payload.items) + ) { + return undefined; + } + return { message, payload, items: payload.items }; +} + +function parseToolCallId(toolCallId: string): ToolCallIdentity | undefined { + if (toolCallId.length === 0) return undefined; + + const separator = toolCallId.indexOf("|"); + if (separator === -1) return { callId: toolCallId }; + if ( + separator === 0 || + separator === toolCallId.length - 1 || + toolCallId.indexOf("|", separator + 1) !== -1 + ) { + return undefined; + } + return { + callId: toolCallId.slice(0, separator), + itemId: toolCallId.slice(separator + 1), + }; +} + +function findNativeToolCall( + items: unknown[], + identity: ToolCallIdentity, +): ToolCallMatch | undefined { + let match: ToolCallMatch | undefined; + let matchingCallIds = 0; + + for (let index = 0; index < items.length; index += 1) { + const item = items[index]; + if (!isRecord(item)) continue; + const kind = item.type; + if (kind !== "function_call" && kind !== "custom_tool_call") continue; + if (item.call_id !== identity.callId) continue; + matchingCallIds++; + if ( + identity.itemId !== undefined && + typeof item.id === "string" && + item.id.length > 0 && + item.id !== identity.itemId + ) { + continue; + } + if (match !== undefined) return undefined; + match = { index, item, kind }; + } + + // OMP synthesizes block ids when the captured native item has no id. + if ( + match && + (typeof match.item.id !== "string" || match.item.id.length === 0) && + matchingCallIds !== 1 + ) { + return undefined; + } + return match; +} + +export function rewriteNativeToolInput( + message: unknown, + toolCallId: string, + input: Record, +): void { + const envelope = getNativeEnvelope(message); + if (!envelope) return; + + const identity = parseToolCallId(toolCallId); + if (!identity) return; + const match = findNativeToolCall(envelope.items, identity); + if (!match) return; + + const field = match.kind === "function_call" ? "arguments" : "input"; + const nextValue = + match.kind === "function_call" + ? JSON.stringify(input) + : typeof input.input === "string" + ? input.input + : ""; + if (typeof nextValue !== "string" || match.item[field] === nextValue) return; + + const items = envelope.items.slice(); + items[match.index] = { ...match.item, [field]: nextValue }; + envelope.message.providerPayload = { ...envelope.payload, items }; +} + +function hasRedactedThinkingContent(message: Record): boolean { + if (!Array.isArray(message.content)) return false; + for (const part of message.content) { + if (isRecord(part) && part.type === "thinking" && part.redacted === true) + return true; + } + return false; +} + +export function canClearNativeReasoning(model: unknown): boolean { + if ( + !isRecord(model) || + model.api !== "openai-codex-responses" || + !isRecord(model.compat) + ) { + return false; + } + const compat = model.compat; + if ( + compat.requiresReasoningContentForAllAssistantTurns !== false || + compat.requiresReasoningContentForToolCalls !== false + ) { + return false; + } + if (compat.whenThinking == null) return true; + return ( + isRecord(compat.whenThinking) && + compat.whenThinking.requiresReasoningContentForAllAssistantTurns === + false && + compat.whenThinking.requiresReasoningContentForToolCalls === false + ); +} + +export function clearNativeReasoning( + message: unknown, + allowed: boolean, +): "not-native" | "cleared" | "preserved" { + const envelope = getNativeEnvelope(message); + if (!envelope) return "not-native"; + if (hasRedactedThinkingContent(envelope.message)) return "preserved"; + + let reasoningCount = 0; + let hasComputerCall = false; + for (const item of envelope.items) { + if (!isRecord(item)) continue; + if (item.type === "computer_call") { + hasComputerCall = true; + continue; + } + if (item.type !== "reasoning") continue; + + reasoningCount += 1; + if ( + typeof item.encrypted_content !== "string" || + item.encrypted_content.length === 0 || + (item.content !== undefined && + (!Array.isArray(item.content) || item.content.length > 0)) + ) { + return "preserved"; + } + } + + if (reasoningCount === 0) return "not-native"; + if (!allowed || envelope.payload.dt !== true || hasComputerCall) { + return "preserved"; + } + + const items = envelope.items.filter( + (item) => !isRecord(item) || item.type !== "reasoning", + ); + envelope.message.providerPayload = { ...envelope.payload, items }; + return "cleared"; +} diff --git a/packages/pi-plugin/src/reasoning-replay-pi.test.ts b/packages/pi-plugin/src/reasoning-replay-pi.test.ts index f0f443da5..ef5e9114a 100644 --- a/packages/pi-plugin/src/reasoning-replay-pi.test.ts +++ b/packages/pi-plugin/src/reasoning-replay-pi.test.ts @@ -128,6 +128,119 @@ describe("clearOldReasoningPi", () => { }); }); + it("clears native-only reasoning and replays the same durable watermark", () => { + const db = makeDb(); + const sessionId = "ses-native-reasoning"; + try { + const native = { + type: "openaiResponsesHistory", + dt: true, + items: [ + { type: "reasoning", encrypted_content: "persisted-only-reasoning" }, + { + type: "function_call", + id: "fc1", + call_id: "call1", + name: "read", + arguments: "{}", + }, + ], + }; + const original = { + role: "assistant", + timestamp: 1, + content: [ + { type: "thinking", thinking: "" }, + { type: "toolCall", id: "call1|fc1", name: "read", arguments: {} }, + ], + providerPayload: native, + }; + const first = structuredClone(original); + const messageIdToMaxTag = new Map(); + messageIdToMaxTag.set("a", 1); + messageIdToMaxTag.set("recent", 10); + const options = { + messageIdToMaxTag, + piMessageStableId: () => "a", + nativeReasoningMayClear: true, + }; + const cleared = clearOldReasoningPi({ + ...options, + messages: [first], + clearReasoningAge: 3, + }); + expect(cleared).toEqual({ cleared: 1, newWatermark: 1 }); + expect(first.providerPayload.items).toEqual([native.items[1]]); + + getOrCreateSessionMeta(db, sessionId); + updateSessionMeta(db, sessionId, { + clearedReasoningThroughTag: cleared.newWatermark, + }); + const resumed = structuredClone(original); + expect( + replayClearedReasoningPi({ + ...options, + messages: [resumed], + db, + sessionId, + }), + ).toBe(1); + expect(resumed.providerPayload).toEqual(first.providerPayload); + expect(native.items[0]).toEqual({ + type: "reasoning", + encrypted_content: "persisted-only-reasoning", + }); + } finally { + db.close(); + } + }); + + it("does not claim to clear reasoning owned by a full native snapshot", () => { + const db = makeDb(); + const sessionId = "ses-native-snapshot"; + try { + const message = { + role: "assistant", + content: [ + { + type: "thinking", + thinking: "required history", + thinkingSignature: "sig", + }, + ], + providerPayload: { + type: "openaiResponsesHistory", + dt: false, + items: [ + { type: "reasoning", encrypted_content: "snapshot-reasoning" }, + ], + }, + }; + const before = JSON.stringify(message); + const messageIdToMaxTag = new Map(); + messageIdToMaxTag.set("a", 1); + messageIdToMaxTag.set("recent", 10); + const options = { + messages: [message], + messageIdToMaxTag, + nativeReasoningMayClear: true, + piMessageStableId: () => "a", + }; + expect(clearOldReasoningPi({ ...options, clearReasoningAge: 3 })).toEqual( + { + cleared: 0, + newWatermark: 0, + }, + ); + getOrCreateSessionMeta(db, sessionId); + updateSessionMeta(db, sessionId, { clearedReasoningThroughTag: 1 }); + expect(replayClearedReasoningPi({ ...options, db, sessionId })).toBe(0); + expect(JSON.stringify(message)).toBe(before); + } finally { + db.close(); + } + }); + it("does nothing when ageCutoff is 0 or below", () => { const messages = [ { diff --git a/packages/pi-plugin/src/reasoning-replay-pi.ts b/packages/pi-plugin/src/reasoning-replay-pi.ts index f30af2cfd..f6e4e1a96 100644 --- a/packages/pi-plugin/src/reasoning-replay-pi.ts +++ b/packages/pi-plugin/src/reasoning-replay-pi.ts @@ -37,6 +37,7 @@ import type { ContextDatabase } from "@magic-context/core/features/magic-context/storage"; import { getOrCreateSessionMeta } from "@magic-context/core/features/magic-context/storage"; import type { TagTarget } from "@magic-context/core/hooks/magic-context/tag-messages"; +import { clearNativeReasoning } from "./native-replay-pi"; type PiTextContent = { type: "text"; text: string }; type PiThinkingContent = { @@ -120,6 +121,7 @@ export function clearOldReasoningPi(args: { messages: unknown[]; messageIdToMaxTag: Map; clearReasoningAge: number; + nativeReasoningMayClear?: boolean; piMessageStableId: (msg: unknown, index: number) => string | undefined; }): { cleared: number; newWatermark: number } { const { messages, messageIdToMaxTag, clearReasoningAge, piMessageStableId } = @@ -146,6 +148,13 @@ export function clearOldReasoningPi(args: { const msgTag = messageIdToMaxTag.get(id) ?? 0; if (msgTag === 0 || msgTag > ageCutoff) continue; + const nativeReasoning = clearNativeReasoning( + msg, + args.nativeReasoningMayClear === true, + ); + if (nativeReasoning === "preserved") continue; + const clearedBefore = cleared; + for (const part of msg.content) { if ( part && @@ -173,6 +182,7 @@ export function clearOldReasoningPi(args: { } } } + if (nativeReasoning === "cleared" && cleared === clearedBefore) cleared++; if (cleared > 0 && msgTag > newWatermark) newWatermark = msgTag; } @@ -251,6 +261,7 @@ export function replayClearedReasoningPi(args: { sessionId: string; messages: unknown[]; messageIdToMaxTag: Map; + nativeReasoningMayClear?: boolean; piMessageStableId: (msg: unknown, index: number) => string | undefined; }): number { const { db, sessionId, messages, messageIdToMaxTag, piMessageStableId } = @@ -272,6 +283,13 @@ export function replayClearedReasoningPi(args: { const msgTag = messageIdToMaxTag.get(id) ?? 0; if (msgTag === 0 || msgTag > watermark) continue; + const nativeReasoning = clearNativeReasoning( + msg, + args.nativeReasoningMayClear === true, + ); + if (nativeReasoning === "preserved") continue; + const clearedBefore = cleared; + for (const part of msg.content) { if ( part && @@ -293,6 +311,7 @@ export function replayClearedReasoningPi(args: { } } } + if (nativeReasoning === "cleared" && cleared === clearedBefore) cleared++; } return cleared; } diff --git a/packages/pi-plugin/src/signal-peek-drain.test.ts b/packages/pi-plugin/src/signal-peek-drain.test.ts index c1c9cd5f4..7fc086bc8 100644 --- a/packages/pi-plugin/src/signal-peek-drain.test.ts +++ b/packages/pi-plugin/src/signal-peek-drain.test.ts @@ -207,12 +207,6 @@ describe("source contract: peek-then-drain in runPipeline (history)", () => { ); }); - test("inline thinking stripping shares the reasoning watermark", () => { - expect(code).toContain("stripInlineThinkingPi({"); - expect(code).toContain("const combinedWatermark = Math.max("); - expect(code).toContain("clearedReasoningThroughTag: combinedWatermark"); - }); - test("model switch reset clears usage, reasoning, failure, limit, and recovery state", () => { expect(code).toContain("clearedReasoningThroughTag: 0"); expect(code).toContain("clearHistorianFailureState(options.db, sessionId)"); diff --git a/packages/pi-plugin/src/transcript-pi.test.ts b/packages/pi-plugin/src/transcript-pi.test.ts index 40569bc64..524cf1547 100644 --- a/packages/pi-plugin/src/transcript-pi.test.ts +++ b/packages/pi-plugin/src/transcript-pi.test.ts @@ -75,6 +75,91 @@ describe("createPiTranscript", () => { expect(transcript.getOutputMessages()).toBe(messages); }); + it("keeps native text stable while reducing a paired tool input", () => { + const db = createTestDb(); + try { + const text = " Keep the indentation and the explanation"; + const reasoning = { + type: "reasoning", + encrypted_content: "unrelated-encrypted-reasoning", + }; + const hostedOutput = { + type: "image_generation_call", + result: "image-data", + }; + const native = { + type: "openaiResponsesHistory", + dt: true, + items: [ + reasoning, + hostedOutput, + { + type: "message", + role: "assistant", + id: "msg-answer", + content: [{ type: "output_text", text }], + }, + { + type: "function_call", + id: "fc-write", + call_id: "call-write", + name: "write", + arguments: JSON.stringify({ content: "large original input" }), + }, + ], + }; + const original = { + ...assistantMessage(text, 1, { + content: [ + { type: "text", text, textSignature: "msg-answer" }, + { + type: "toolCall", + id: "call-write|fc-write", + name: "write", + arguments: { content: "large original input" }, + }, + ], + }), + providerPayload: native, + }; + const messages = [ + original, + toolResultMessage("call-write|fc-write", "written", 2), + ]; + const sessionId = "ses-native-replay"; + const tagger = createTagger(); + tagger.initFromDb(sessionId, db); + const tagged = createPiTranscript(messages, sessionId); + tagTranscript(sessionId, tagged, tagger, db); + tagged.commit(); + const taggedAssistant = messages[0] as typeof original; + expect(textOf(taggedAssistant)).not.toBe(text); + expect(taggedAssistant.providerPayload).toBe(native); + + const reduced = createPiTranscript(messages, sessionId); + const toolPart = reduced.messages[0]?.parts[1]; + if (!toolPart) throw new Error("missing tool part"); + toolPart.replaceWithSentinel("[dropped input]"); + reduced.commit(); + const payload = (messages[0] as typeof original).providerPayload; + expect(payload.items[2]).toEqual(native.items[2]); + expect(payload.items[3]).toMatchObject({ + type: "function_call", + id: "fc-write", + call_id: "call-write", + name: "write", + arguments: JSON.stringify({ dropped: "[dropped input]" }), + }); + expect(payload.items[0]).toEqual(reasoning); + expect(payload.items[1]).toEqual(hostedOutput); + expect(native.items[3]).toMatchObject({ + arguments: JSON.stringify({ content: "large original input" }), + }); + } finally { + closeQuietly(db); + } + }); + it("leaves Pi and OMP whitespace-only assistant framing untagged after peeling MC tags", () => { const vectors = [ { diff --git a/packages/pi-plugin/src/transcript-pi.ts b/packages/pi-plugin/src/transcript-pi.ts index 5b62df5a6..edec56fa5 100644 --- a/packages/pi-plugin/src/transcript-pi.ts +++ b/packages/pi-plugin/src/transcript-pi.ts @@ -77,6 +77,7 @@ import type { TranscriptPart, TranscriptPartKind, } from "@magic-context/core/shared/transcript"; +import { rewriteNativeToolInput } from "./native-replay-pi"; import { resolvePiHarnessKind } from "./pi-harness-kind"; import { resolvePiStableId, SYNTH_USER_ID_PREFIX } from "./read-session-pi"; @@ -674,6 +675,7 @@ function createPiAssistantPart( ...(working[messageIndex] as PiAssistantMessage), content: newContent, }; + rewriteNativeToolInput(working[messageIndex], p.id, replacementArgs); markDirty(messageIndex); return true; } @@ -733,6 +735,7 @@ function createPiAssistantPart( ...(working[messageIndex] as PiAssistantMessage), content: newContent, }; + rewriteNativeToolInput(working[messageIndex], p.id, input); markDirty(messageIndex); return true; }, @@ -777,6 +780,14 @@ function createPiAssistantPart( ...(working[messageIndex] as PiAssistantMessage), content: newContent, }; + const replacement = newContent[partIndex]; + if (existing?.type === "toolCall" && replacement?.type === "toolCall") { + rewriteNativeToolInput( + working[messageIndex], + existing.id, + replacement.arguments, + ); + } markDirty(messageIndex); return true; }, From b99a3fc37851f7e8d72c7dee497ac72e31d2db06 Mon Sep 17 00:00:00 2001 From: SilentBless <126187696+SilentBless@users.noreply.github.com> Date: Thu, 10 Sep 2026 09:55:14 +0300 Subject: [PATCH 2/7] fix(pi): preserve local cleanup and native serialization safety --- packages/pi-plugin/README.md | 2 + .../pi-plugin/src/context-handler.test.ts | 70 +++++ packages/pi-plugin/src/context-handler.ts | 8 +- .../pi-plugin/src/native-replay-pi.test.ts | 56 ++++ packages/pi-plugin/src/native-replay-pi.ts | 16 +- .../pi-plugin/src/reasoning-replay-pi.test.ts | 253 ++++++++++++++---- packages/pi-plugin/src/reasoning-replay-pi.ts | 6 +- 7 files changed, 354 insertions(+), 57 deletions(-) diff --git a/packages/pi-plugin/README.md b/packages/pi-plugin/README.md index f6a7c35da..a6fb70bb5 100644 --- a/packages/pi-plugin/README.md +++ b/packages/pi-plugin/README.md @@ -78,6 +78,8 @@ OMP's legacy Pi loader maps `@earendil-works/*` imports to its bundled `@oh-my-p OMP can replay Responses/Codex history from `providerPayload` instead of ordinary message content. Tool-input reductions update the matched native call without discarding unrelated native results. Old encrypted reasoning is cleared only for Codex models whose resolved compatibility settings explicitly allow omission, and only from incremental history. Other models, full snapshots, plaintext or malformed reasoning, redacted thinking, and computer-linked reasoning retain their native reasoning. Native text and user/developer history carriers are not rewritten. +Display summaries (`summary`) are removed with eligible old encrypted reasoning; they are the source of ordinary Pi `thinking`, not a native preservation requirement. Retaining a native payload does not prevent per-part cleanup of stale non-redacted Pi thinking and its signature. + --- ## Configuration diff --git a/packages/pi-plugin/src/context-handler.test.ts b/packages/pi-plugin/src/context-handler.test.ts index 59c2a630a..a20dfc3aa 100644 --- a/packages/pi-plugin/src/context-handler.test.ts +++ b/packages/pi-plugin/src/context-handler.test.ts @@ -3619,6 +3619,76 @@ describe("registerPiContextHandler", () => { } }); + it("replays an inline-only watermark after restart without fresh age cleanup", async () => { + const db = createTestDb(); + const sessionId = "ses-inline-reasoning-watermark"; + try { + const fake = createFakePi(); + registerPiContextHandler(fake.pi as never, { + db, + heuristics: { clearReasoningAge: 1 }, + scheduler: { executeThresholdPercentage: 80 }, + }); + let handler = fake.handlers.get("context") as ( + event: { messages: never[] }, + ctx: never, + ) => Promise<{ messages: never[] } | undefined>; + const runPass = async (percent: number, newUser = false) => { + const messages = [ + userMessage("first", 1), + assistantMessage( + "Keep stale private thought visible", + 2, + ), + userMessage("second", 3), + assistantMessage("latest answer", 4), + ]; + const entryIds = [ + "entry-u1", + "entry-inline", + "entry-u2", + "entry-latest", + ]; + if (newUser) { + messages.push(userMessage("new request", 5)); + entryIds.push("entry-new"); + } + const result = await handler({ messages: messages as never[] }, { + ...fakeContext(sessionId, process.cwd(), entryIds, messages as never), + getContextUsage: () => ({ + tokens: percent * 1_000, + percent, + contextWindow: 100_000, + }), + } as never); + if (!result) throw new Error("expected transformed messages"); + return result.messages; + }; + + const executed = await runPass(90); + expect(textOf(executed[1])).toContain("Keep visible"); + expect(textOf(executed[1])).not.toContain("stale private thought"); + updateSessionMeta(db, sessionId, { + lastResponseTime: Date.now(), + cacheTtl: "59m", + lastContextPercentage: 1, + lastInputTokens: 1_000, + }); + clearContextHandlerSession(sessionId); + registerPiContextHandler(fake.pi as never, { + db, + heuristics: { clearReasoningAge: 100 }, + scheduler: { executeThresholdPercentage: 80 }, + }); + handler = fake.handlers.get("context") as typeof handler; + const replayed = await runPass(1, true); + expect(textOf(replayed[1])).toBe(textOf(executed[1])); + } finally { + clearContextHandlerSession(sessionId); + closeQuietly(db); + } + }); + it("restores reasoning bytes when the durable watermark write fails", async () => { const db = createTestDb(); const sessionId = "ses-reasoning-watermark-failure"; diff --git a/packages/pi-plugin/src/context-handler.ts b/packages/pi-plugin/src/context-handler.ts index 27088f6d2..7999c5f15 100644 --- a/packages/pi-plugin/src/context-handler.ts +++ b/packages/pi-plugin/src/context-handler.ts @@ -4624,6 +4624,7 @@ function pendingPiMarkerCoveredByRenderedBoundary( function captureReasoningMutationRollback( messages: readonly unknown[], + nativeReasoningMayClear: boolean, ): () => void { const snapshots: Array<{ part: Record; @@ -4641,7 +4642,7 @@ function captureReasoningMutationRollback( }; if (message.role !== "assistant" || !Array.isArray(message.content)) continue; - if (Object.hasOwn(message, "providerPayload")) { + if (nativeReasoningMayClear && Object.hasOwn(message, "providerPayload")) { snapshots.push({ part: message, field: "providerPayload", @@ -5658,7 +5659,10 @@ async function runPipeline(args: RunPipelineArgs): Promise { // wire on a pass that already dropped tools (inconsistent + a missed // same-pass mutation). shouldRunHeuristics is the broader, correct set. if (args.reasoningClearing && shouldRunHeuristics && routineCleanupApplied) { - const rollbackReasoning = captureReasoningMutationRollback(workingMessages); + const rollbackReasoning = captureReasoningMutationRollback( + workingMessages, + args.reasoningClearing.nativeReasoningMayClear, + ); try { const tClearReasoning = performance.now(); const prevWatermark = args.sessionMeta.clearedReasoningThroughTag ?? 0; diff --git a/packages/pi-plugin/src/native-replay-pi.test.ts b/packages/pi-plugin/src/native-replay-pi.test.ts index a511cb1b0..905869b60 100644 --- a/packages/pi-plugin/src/native-replay-pi.test.ts +++ b/packages/pi-plugin/src/native-replay-pi.test.ts @@ -84,6 +84,43 @@ describe("native tool input reductions", () => { expect(call.input).toBe("original patch"); }); + it("leaves native function input untouched when generic arguments cannot produce JSON", () => { + const circular: Record = {}; + circular.self = circular; + const unsupportedInputs: Record[] = [ + circular, + { count: BigInt(1) }, + { + toJSON() { + throw new Error("serialization failed"); + }, + }, + { + toJSON() { + return undefined; + }, + }, + ]; + + for (const input of unsupportedInputs) { + const call = { + type: "function_call", + id: "fc1", + call_id: "call1", + name: "read", + arguments: JSON.stringify({ path: "original.txt" }), + }; + const message = nativeMessage([call]); + const payload = message.providerPayload; + const before = structuredClone(payload); + + rewriteNativeToolInput(message, "call1|fc1", input); + + expect(message.providerPayload).toBe(payload); + expect(message.providerPayload).toEqual(before); + } + }); + it("matches unique id-less calls when OMP supplies a synthesized block id", () => { const call = { type: "function_call", @@ -250,6 +287,25 @@ describe("native reasoning retention", () => { expect(original.items).toEqual([reasoning, text, call, image]); }); + it("clears eligible encrypted native reasoning with a display summary", () => { + const reasoning = { + type: "reasoning", + encrypted_content: "opaque-encrypted-content", + summary: [{ type: "summary_text", text: "displayed thinking" }], + }; + const text = { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "keep" }], + }; + const message = nativeMessage([reasoning, text]); + const original = message.providerPayload; + + expect(clearNativeReasoning(message, true)).toBe("cleared"); + expect(message.providerPayload.items).toEqual([text]); + expect(original.items).toEqual([reasoning, text]); + }); + it("retains native reasoning when the active model has not authorized omission", () => { const message = nativeMessage([ { type: "reasoning", encrypted_content: "must-stay" }, diff --git a/packages/pi-plugin/src/native-replay-pi.ts b/packages/pi-plugin/src/native-replay-pi.ts index 574d8d9ad..44b7a060c 100644 --- a/packages/pi-plugin/src/native-replay-pi.ts +++ b/packages/pi-plugin/src/native-replay-pi.ts @@ -99,12 +99,16 @@ export function rewriteNativeToolInput( if (!match) return; const field = match.kind === "function_call" ? "arguments" : "input"; - const nextValue = - match.kind === "function_call" - ? JSON.stringify(input) - : typeof input.input === "string" - ? input.input - : ""; + let nextValue: string | undefined; + if (match.kind === "function_call") { + try { + nextValue = JSON.stringify(input); + } catch { + return; + } + } else { + nextValue = typeof input.input === "string" ? input.input : ""; + } if (typeof nextValue !== "string" || match.item[field] === nextValue) return; const items = envelope.items.slice(); diff --git a/packages/pi-plugin/src/reasoning-replay-pi.test.ts b/packages/pi-plugin/src/reasoning-replay-pi.test.ts index ef5e9114a..344c80126 100644 --- a/packages/pi-plugin/src/reasoning-replay-pi.test.ts +++ b/packages/pi-plugin/src/reasoning-replay-pi.test.ts @@ -195,11 +195,16 @@ describe("clearOldReasoningPi", () => { } }); - it("does not claim to clear reasoning owned by a full native snapshot", () => { + it("preserves a full native snapshot while clearing local Pi thinking and replaying its watermark", () => { const db = makeDb(); const sessionId = "ses-native-snapshot"; try { - const message = { + const native = { + type: "openaiResponsesHistory", + dt: false, + items: [{ type: "reasoning", encrypted_content: "snapshot-reasoning" }], + }; + const original = { role: "assistant", content: [ { @@ -207,35 +212,146 @@ describe("clearOldReasoningPi", () => { thinking: "required history", thinkingSignature: "sig", }, + { type: "text", text: "reply" }, ], - providerPayload: { - type: "openaiResponsesHistory", - dt: false, - items: [ - { type: "reasoning", encrypted_content: "snapshot-reasoning" }, - ], - }, + providerPayload: native, }; - const before = JSON.stringify(message); - const messageIdToMaxTag = new Map(); - messageIdToMaxTag.set("a", 1); - messageIdToMaxTag.set("recent", 10); + const nativeSnapshot = structuredClone(native); + const messageIdToMaxTag = new Map([ + ["a", 1], + ["recent", 10], + ]); const options = { - messages: [message], messageIdToMaxTag, nativeReasoningMayClear: true, piMessageStableId: () => "a", }; - expect(clearOldReasoningPi({ ...options, clearReasoningAge: 3 })).toEqual( + + const first = structuredClone(original); + const cleared = clearOldReasoningPi({ + ...options, + messages: [first], + clearReasoningAge: 3, + }); + expect(cleared).toEqual({ cleared: 1, newWatermark: 1 }); + expect(first.content).toEqual([ + { type: "thinking", thinking: "" }, + { type: "text", text: "reply" }, + ]); + expect(first.providerPayload).toEqual(nativeSnapshot); + + getOrCreateSessionMeta(db, sessionId); + updateSessionMeta(db, sessionId, { + clearedReasoningThroughTag: cleared.newWatermark, + }); + const resumed = structuredClone(original); + expect( + replayClearedReasoningPi({ + ...options, + messages: [resumed], + db, + sessionId, + }), + ).toBe(1); + expect(resumed).toEqual(first); + expect(resumed.providerPayload).toEqual(nativeSnapshot); + } finally { + db.close(); + } + }); + + it("cleans local thinking when native clearing is disallowed or native reasoning has plaintext", () => { + const db = makeDb(); + try { + for (const fixture of [ { - cleared: 0, - newWatermark: 0, + label: "native-clearing-disallowed", + nativeReasoningMayClear: false, + providerPayload: { + type: "openaiResponsesHistory", + dt: true, + items: [ + { + type: "reasoning", + encrypted_content: "opaque-encrypted-content", + }, + ], + }, }, - ); - getOrCreateSessionMeta(db, sessionId); - updateSessionMeta(db, sessionId, { clearedReasoningThroughTag: 1 }); - expect(replayClearedReasoningPi({ ...options, db, sessionId })).toBe(0); - expect(JSON.stringify(message)).toBe(before); + { + label: "plaintext-native-reasoning", + nativeReasoningMayClear: true, + providerPayload: { + type: "openaiResponsesHistory", + dt: true, + items: [ + { + type: "reasoning", + encrypted_content: "opaque-encrypted-content", + content: [ + { + type: "reasoning_text", + text: "required plaintext", + }, + ], + }, + ], + }, + }, + ] as const) { + const original = { + role: "assistant", + content: [ + { + type: "thinking", + thinking: "ordinary local thinking", + thinkingSignature: "local-signature", + }, + { type: "text", text: "reply" }, + ], + providerPayload: fixture.providerPayload, + }; + const nativeSnapshot = structuredClone(fixture.providerPayload); + const messageIdToMaxTag = new Map([ + ["a", 1], + ["recent", 4], + ]); + const options = { + messageIdToMaxTag, + nativeReasoningMayClear: fixture.nativeReasoningMayClear, + piMessageStableId: () => "a", + }; + + const first = structuredClone(original); + const cleared = clearOldReasoningPi({ + ...options, + messages: [first], + clearReasoningAge: 2, + }); + expect(cleared).toEqual({ cleared: 1, newWatermark: 1 }); + expect(first.content).toEqual([ + { type: "thinking", thinking: "" }, + { type: "text", text: "reply" }, + ]); + expect(first.providerPayload).toEqual(nativeSnapshot); + + const sessionId = `ses-${fixture.label}`; + getOrCreateSessionMeta(db, sessionId); + updateSessionMeta(db, sessionId, { + clearedReasoningThroughTag: cleared.newWatermark, + }); + const resumed = structuredClone(original); + expect( + replayClearedReasoningPi({ + ...options, + messages: [resumed], + db, + sessionId, + }), + ).toBe(1); + expect(resumed).toEqual(first); + expect(resumed.providerPayload).toEqual(nativeSnapshot); + } } finally { db.close(); } @@ -262,15 +378,23 @@ describe("clearOldReasoningPi", () => { expect(result.newWatermark).toBe(0); }); - it("leaves REDACTED thinking blocks untouched (signature + data preserved)", () => { - // A redacted block bypasses the empty-thinking drop in Pi's serializers - // (transform-messages.ts / anthropic.ts serialize `redacted` before the - // empty check), so emptying it + dropping the signature would leave a - // malformed redacted block on the wire. It must be preserved verbatim. - const messages = [ - { + it("preserves redacted and native blocks while clearing ordinary local thinking", () => { + // Redacted blocks serialize before empty blocks, so they remain verbatim; + // ordinary local thinking in the same message must still be cleared. + const db = makeDb(); + const sessionId = "ses-mixed-generic-thinking"; + try { + const native = { + type: "openaiResponsesHistory", + dt: true, + items: [ + { type: "reasoning", encrypted_content: "opaque native reasoning" }, + ], + }; + const original = { role: "assistant", timestamp: 1, + providerPayload: native, content: [ { type: "thinking", @@ -278,25 +402,60 @@ describe("clearOldReasoningPi", () => { thinkingSignature: "sig-abc", redacted: true, }, + { + type: "thinking", + thinking: "ordinary local thinking", + thinkingSignature: "ordinary-signature", + }, + { type: "text", text: "reply" }, ], - }, - ]; - const id0 = piMessageStableId(messages[0], 0); - if (!id0) throw new Error("piMessageStableId returned undefined"); - const messageIdToMaxTag = new Map([[id0, 1]]); - const result = clearOldReasoningPi({ - messages, - messageIdToMaxTag, - clearReasoningAge: 1, - piMessageStableId, - }); - expect(result.cleared).toBe(0); - expect(messages[0].content[0]).toMatchObject({ - type: "thinking", - thinking: "opaque-redacted-payload", - thinkingSignature: "sig-abc", - redacted: true, - }); + }; + const messageIdToMaxTag = new Map([ + [requireId(original, 0), 1], + ["recent", 4], + ]); + const options = { + messageIdToMaxTag, + piMessageStableId, + nativeReasoningMayClear: true, + }; + + const first = structuredClone(original); + const cleared = clearOldReasoningPi({ + ...options, + messages: [first], + clearReasoningAge: 2, + }); + expect(cleared).toEqual({ cleared: 1, newWatermark: 1 }); + expect(first.content).toEqual([ + { + type: "thinking", + thinking: "opaque-redacted-payload", + thinkingSignature: "sig-abc", + redacted: true, + }, + { type: "thinking", thinking: "" }, + { type: "text", text: "reply" }, + ]); + expect(first.providerPayload).toEqual(native); + + getOrCreateSessionMeta(db, sessionId); + updateSessionMeta(db, sessionId, { + clearedReasoningThroughTag: cleared.newWatermark, + }); + const resumed = structuredClone(original); + expect( + replayClearedReasoningPi({ + ...options, + messages: [resumed], + db, + sessionId, + }), + ).toBe(1); + expect(resumed).toEqual(first); + } finally { + db.close(); + } }); }); diff --git a/packages/pi-plugin/src/reasoning-replay-pi.ts b/packages/pi-plugin/src/reasoning-replay-pi.ts index f6e4e1a96..c65ec4f7f 100644 --- a/packages/pi-plugin/src/reasoning-replay-pi.ts +++ b/packages/pi-plugin/src/reasoning-replay-pi.ts @@ -152,7 +152,8 @@ export function clearOldReasoningPi(args: { msg, args.nativeReasoningMayClear === true, ); - if (nativeReasoning === "preserved") continue; + // Native preservation protects providerPayload only; local Pi thinking + // still follows the ordinary per-part cleanup policy. const clearedBefore = cleared; for (const part of msg.content) { @@ -287,7 +288,8 @@ export function replayClearedReasoningPi(args: { msg, args.nativeReasoningMayClear === true, ); - if (nativeReasoning === "preserved") continue; + // Native preservation protects providerPayload only; local Pi thinking + // still follows the ordinary per-part cleanup policy. const clearedBefore = cleared; for (const part of msg.content) { From 10d797ccd0591a425aca83e1ecc9e678ce3fe184 Mon Sep 17 00:00:00 2001 From: SilentBless <126187696+SilentBless@users.noreply.github.com> Date: Thu, 10 Sep 2026 22:00:45 +0300 Subject: [PATCH 3/7] docs(pi): clarify native history mutation boundary --- packages/pi-plugin/README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/packages/pi-plugin/README.md b/packages/pi-plugin/README.md index a6fb70bb5..855790d31 100644 --- a/packages/pi-plugin/README.md +++ b/packages/pi-plugin/README.md @@ -80,6 +80,8 @@ OMP can replay Responses/Codex history from `providerPayload` instead of ordinar Display summaries (`summary`) are removed with eligible old encrypted reasoning; they are the source of ordinary Pi `thinking`, not a native preservation requirement. Retaining a native payload does not prevent per-part cleanup of stale non-redacted Pi thinking and its signature. +OMP exposes no dedicated native-item mutation API, so this adapter targets its current Responses history representation. It changes the request transcript; it does not perform a stored session-JSONL rewrite. + --- ## Configuration From f16139294d46bc4c4e69cc70459ac7db6307e272 Mon Sep 17 00:00:00 2001 From: SilentBless <126187696+SilentBless@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:32:09 +0300 Subject: [PATCH 4/7] fix(pi): persist authorized native replay activation --- packages/pi-plugin/README.md | 4 + .../pi-plugin/src/clone-inheritance.test.ts | 65 +++ packages/pi-plugin/src/context-handler.ts | 52 ++- .../src/native-replay-state-pi.test.ts | 411 ++++++++++++++++++ .../pi-plugin/src/native-replay-state-pi.ts | 157 +++++++ .../pi-plugin/src/reasoning-replay-pi.test.ts | 208 ++------- packages/pi-plugin/src/reasoning-replay-pi.ts | 43 +- packages/pi-plugin/src/transcript-pi.test.ts | 201 +++++---- packages/pi-plugin/src/transcript-pi.ts | 62 +-- packages/plugin/scripts/clone-session.test.ts | 39 ++ packages/plugin/scripts/clone-session.ts | 7 +- .../features/magic-context/storage-clone.ts | 60 +++ .../src/features/magic-context/storage-db.ts | 2 + .../storage-native-replay.test.ts | 185 ++++++++ .../magic-context/storage-native-replay.ts | 185 ++++++++ 15 files changed, 1367 insertions(+), 314 deletions(-) create mode 100644 packages/pi-plugin/src/native-replay-state-pi.test.ts create mode 100644 packages/pi-plugin/src/native-replay-state-pi.ts create mode 100644 packages/plugin/src/features/magic-context/storage-native-replay.test.ts create mode 100644 packages/plugin/src/features/magic-context/storage-native-replay.ts diff --git a/packages/pi-plugin/README.md b/packages/pi-plugin/README.md index 855790d31..89f7a338f 100644 --- a/packages/pi-plugin/README.md +++ b/packages/pi-plugin/README.md @@ -78,6 +78,10 @@ OMP's legacy Pi loader maps `@earendil-works/*` imports to its bundled `@oh-my-p OMP can replay Responses/Codex history from `providerPayload` instead of ordinary message content. Tool-input reductions update the matched native call without discarding unrelated native results. Old encrypted reasoning is cleared only for Codex models whose resolved compatibility settings explicitly allow omission, and only from incremental history. Other models, full snapshots, plaintext or malformed reasoning, redacted thinking, and computer-linked reasoning retain their native reasoning. Native text and user/developer history carriers are not rewritten. +Native tool-input values and native reasoning removals have separate persisted replay state. Existing dropped tags and local reasoning watermarks do not activate native changes during upgrade: first application waits for an already-authorized cache-busting pass and is persisted before publication. Deferred passes replay only those saved native decisions. + +Function calls carry the canonical dropped-marker JSON. Custom calls use OMP's existing empty-string fallback, which does not carry that marker or imply the same copied-input rejection behavior. + Display summaries (`summary`) are removed with eligible old encrypted reasoning; they are the source of ordinary Pi `thinking`, not a native preservation requirement. Retaining a native payload does not prevent per-part cleanup of stale non-redacted Pi thinking and its signature. OMP exposes no dedicated native-item mutation API, so this adapter targets its current Responses history representation. It changes the request transcript; it does not perform a stored session-JSONL rewrite. diff --git a/packages/pi-plugin/src/clone-inheritance.test.ts b/packages/pi-plugin/src/clone-inheritance.test.ts index 3e6206319..43872850b 100644 --- a/packages/pi-plugin/src/clone-inheritance.test.ts +++ b/packages/pi-plugin/src/clone-inheritance.test.ts @@ -11,6 +11,12 @@ import { getSourceContents, getTagsBySession, } from "@magic-context/core/features/magic-context/storage"; +import { + addNativeReasoningIds, + getNativeReasoningIds, + getNativeToolInputs, + saveNativeToolInputs, +} from "@magic-context/core/features/magic-context/storage-native-replay"; import { replayCavemanCompression } from "@magic-context/core/hooks/magic-context/caveman-cleanup"; import type { TagTarget } from "@magic-context/core/hooks/magic-context/tag-messages"; import type { Database } from "@magic-context/core/shared/sqlite"; @@ -648,6 +654,65 @@ describe("Pi clone state inheritance", () => { expect(JSON.parse(row.images)).toEqual(["u1"]); }); + it("carries only retained native tool and reasoning decisions into a remapped clone", () => { + const database = db(); + const retainedInput = + '{"path":"src/retained.ts","range":{"start":10,"end":40},"marker":"[truncated]"}'; + seedTag(database, { + tagNumber: 1, + messageId: "call-retained", + type: "tool", + ownerId: "assistant-retained", + }); + seedTag(database, { + tagNumber: 2, + messageId: "call-outside", + type: "tool", + ownerId: "assistant-outside", + }); + saveNativeToolInputs( + database, + "source", + new Map([ + ["call-retained", retainedInput], + ["call-outside", '{"path":"src/outside.ts"}'], + ]), + ); + addNativeReasoningIds(database, "source", [ + "assistant-retained", + "assistant-outside", + ]); + + const filter: CloneSessionStateFilter = { + resolveBoundaryOrdinal: () => undefined, + includeTag: (tag) => + tag.type === "tool" && tag.toolOwnerMessageId === "assistant-retained", + includeMessageId: (id) => id === "assistant-retained", + mapMessageId: (id) => + id === "assistant-retained" + ? "clone-assistant-retained" + : id === "call-retained" + ? "clone-call-retained" + : id, + selectPendingPiMarker: () => null, + }; + + const result = copySessionStateForClone( + database, + "source", + "clone", + filter, + ); + + expect(result.tagsCopied).toBe(1); + expect(getNativeToolInputs(database, "clone")).toEqual( + new Map([["clone-call-retained", retainedInput]]), + ); + expect(getNativeReasoningIds(database, "clone")).toEqual( + new Set(["clone-assistant-retained"]), + ); + }); + it("leaves every m0/m1 cache field fresh so the first pass hard-materializes", () => { const database = db(); seedCompartment(database, { sequence: 1, startId: "u1", endId: "a1" }); diff --git a/packages/pi-plugin/src/context-handler.ts b/packages/pi-plugin/src/context-handler.ts index 7999c5f15..414566201 100644 --- a/packages/pi-plugin/src/context-handler.ts +++ b/packages/pi-plugin/src/context-handler.ts @@ -224,6 +224,10 @@ import { trimPiMessagesToCachedBoundary, } from "./inject-compartments-pi"; import { canClearNativeReasoning } from "./native-replay-pi"; +import { + applyNativeReasoningReplayPi, + applyNativeToolInputReplayPi, +} from "./native-replay-state-pi"; import { hasVisibleNoteReadCallPi } from "./note-visibility-pi"; import { resolvePiUsableContextLimit, @@ -4624,11 +4628,10 @@ function pendingPiMarkerCoveredByRenderedBoundary( function captureReasoningMutationRollback( messages: readonly unknown[], - nativeReasoningMayClear: boolean, ): () => void { const snapshots: Array<{ part: Record; - field: "thinking" | "text" | "providerPayload"; + field: "thinking" | "text"; value: unknown; hadSignature?: boolean; signature?: unknown; @@ -4638,17 +4641,9 @@ function captureReasoningMutationRollback( const message = raw as { role?: unknown; content?: unknown; - providerPayload?: unknown; }; if (message.role !== "assistant" || !Array.isArray(message.content)) continue; - if (nativeReasoningMayClear && Object.hasOwn(message, "providerPayload")) { - snapshots.push({ - part: message, - field: "providerPayload", - value: message.providerPayload, - }); - } for (const rawPart of message.content) { if (!rawPart || typeof rawPart !== "object") continue; const part = rawPart as Record; @@ -5354,7 +5349,6 @@ async function runPipeline(args: RunPipelineArgs): Promise { sessionId: args.sessionId, messages: workingMessages, messageIdToMaxTag, - nativeReasoningMayClear: args.reasoningClearing.nativeReasoningMayClear, piMessageStableId: stableIdResolver, }); const inlineReplay = replayStrippedInlineThinkingPi({ @@ -5658,11 +5652,9 @@ async function runPipeline(args: RunPipelineArgs): Promise { // materialization passes where heuristics DO run — leaving reasoning on the // wire on a pass that already dropped tools (inconsistent + a missed // same-pass mutation). shouldRunHeuristics is the broader, correct set. + let reasoningPersistenceFailed = false; if (args.reasoningClearing && shouldRunHeuristics && routineCleanupApplied) { - const rollbackReasoning = captureReasoningMutationRollback( - workingMessages, - args.reasoningClearing.nativeReasoningMayClear, - ); + const rollbackReasoning = captureReasoningMutationRollback(workingMessages); try { const tClearReasoning = performance.now(); const prevWatermark = args.sessionMeta.clearedReasoningThroughTag ?? 0; @@ -5670,7 +5662,6 @@ async function runPipeline(args: RunPipelineArgs): Promise { messages: workingMessages, messageIdToMaxTag, clearReasoningAge: args.reasoningClearing.clearReasoningAge, - nativeReasoningMayClear: args.reasoningClearing.nativeReasoningMayClear, piMessageStableId: stableIdResolver, }); const stripOutcome = stripInlineThinkingPi({ @@ -5707,6 +5698,7 @@ async function runPipeline(args: RunPipelineArgs): Promise { executedWorkThisPass = true; } } catch (err) { + reasoningPersistenceFailed = true; // Never ship cleared reasoning unless replay state persisted. Restoring the // pre-cleanup parts keeps this pass byte-stable and lets the next execute // pass retry the watermark write. @@ -5836,6 +5828,34 @@ async function runPipeline(args: RunPipelineArgs): Promise { const tTranscriptCommit = performance.now(); transcript.commit(); logTransformTiming(args.sessionId, "transcriptCommit", tTranscriptCommit); + + // Legacy drop/watermark state does not authorize first native activation. + // Use committed canonical inputs, then publish only persisted native decisions. + const nativeInputsApplied = applyNativeToolInputReplayPi({ + db: args.db, + sessionId: args.sessionId, + messages: args.messages, + changes: transcript.getToolInputChanges(), + canApply: isCacheBustingPass, + }); + const nativeReasoningApplied = args.reasoningClearing + ? applyNativeReasoningReplayPi({ + db: args.db, + sessionId: args.sessionId, + messages: args.messages, + messageIdToMaxTag, + stableId: stableIdResolver, + localWatermark: args.sessionMeta.clearedReasoningThroughTag ?? 0, + clearReasoningAge: args.reasoningClearing.clearReasoningAge, + omissionAllowed: args.reasoningClearing.nativeReasoningMayClear, + canApply: isCacheBustingPass && !reasoningPersistenceFailed, + detectAged: shouldRunHeuristics && routineCleanupApplied, + }) + : 0; + if (nativeInputsApplied > 0 || nativeReasoningApplied > 0) { + heuristicOrReasoningDidMutate = true; + executedWorkThisPass = true; + } if (toolReclaimApplicationOpportunity) { advanceToolReclaimWatermarkToCurrentMax(args.db, args.sessionId); } diff --git a/packages/pi-plugin/src/native-replay-state-pi.test.ts b/packages/pi-plugin/src/native-replay-state-pi.test.ts new file mode 100644 index 000000000..dcd7184be --- /dev/null +++ b/packages/pi-plugin/src/native-replay-state-pi.test.ts @@ -0,0 +1,411 @@ +import { describe, expect, it } from "bun:test"; +import { + getOrCreateSessionMeta, + getPendingOps, + getTagsBySession, + updateSessionMeta, + updateTagDropMode, + updateTagStatus, +} from "@magic-context/core/features/magic-context/storage"; +import { + getNativeReasoningIds, + getNativeToolInputs, +} from "@magic-context/core/features/magic-context/storage-native-replay"; +import { + clearContextHandlerSession, + registerPiContextHandler, +} from "./context-handler"; +import { + applyNativeReasoningReplayPi, + applyNativeToolInputReplayPi, +} from "./native-replay-state-pi"; +import { + assistantMessage, + createFakePi, + createTestDb, + fakeContext, + toolResultMessage, + userMessage, +} from "./test-utils.test"; + +const model = { + id: "native-upgrade-codex", + api: "openai-codex-responses", + provider: "openai-codex", + contextWindow: 100_000, + compat: { + requiresReasoningContentForAllAssistantTurns: false, + requiresReasoningContentForToolCalls: false, + }, +}; +const callId = "call-old|fc-old"; +const staleInput = "retained original argument"; +const ciphertext = "retained original ciphertext"; + +function oldAssistant() { + return assistantMessage("visible answer", 2, { + api: model.api, + provider: model.provider, + model: model.id, + content: [ + { + type: "thinking", + thinking: "old summary", + thinkingSignature: "old signature", + }, + { type: "text", text: "visible answer" }, + { + type: "toolCall", + id: callId, + name: "read", + arguments: { path: staleInput }, + }, + ], + providerPayload: { + type: "openaiResponsesHistory", + provider: model.provider, + dt: true, + items: [ + { + type: "reasoning", + encrypted_content: ciphertext, + summary: [{ type: "summary_text", text: "old summary" }], + }, + { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "visible answer" }], + }, + { + type: "function_call", + id: "fc-old", + call_id: "call-old", + name: "read", + arguments: JSON.stringify({ path: staleInput }), + }, + { + type: "image_generation_call", + id: "image-old", + status: "completed", + result: "aW1hZ2U=", + }, + ], + }, + }); +} + +function nativeBytes(messages: unknown[]): string { + return JSON.stringify( + messages + .filter( + (message): message is { providerPayload: unknown } => + !!message && + typeof message === "object" && + "providerPayload" in message, + ) + .map((message) => message.providerPayload), + ); +} + +function fixture(sessionId: string) { + const db = createTestDb(); + const fake = createFakePi(); + const options = { + db, + heuristics: { clearReasoningAge: 100 }, + scheduler: { executeThresholdPercentage: 80 }, + }; + const restart = () => { + clearContextHandlerSession(sessionId); + registerPiContextHandler(fake.pi as never, options); + }; + restart(); + const pass = async (percent: number) => { + updateSessionMeta(db, sessionId, { + lastResponseTime: Date.now(), + cacheTtl: "59m", + }); + const messages = [ + userMessage("first", 1), + oldAssistant(), + toolResultMessage(callId, "old output", 3), + userMessage("continue", 4), + ]; + const handler = fake.handlers.get("context") as ( + event: { messages: never[] }, + ctx: never, + ) => Promise<{ messages: unknown[] } | undefined>; + const result = await handler({ messages: messages as never[] }, { + ...fakeContext( + sessionId, + process.cwd(), + ["entry-user", "entry-old", "entry-result", "entry-next"], + messages, + ), + model, + getContextUsage: () => ({ + percent, + tokens: percent * 1000, + contextWindow: 100_000, + }), + } as never); + if (!result) throw new Error("Missing context result"); + return result.messages; + }; + const seedLegacy = async () => { + await pass(0); + const tags = getTagsBySession(db, sessionId); + const tool = tags.find( + (tag) => tag.type === "tool" && tag.messageId === callId, + ); + if (!tool) throw new Error("Missing legacy tool tag"); + updateTagStatus(db, sessionId, tool.tagNumber, "dropped"); + updateTagDropMode(db, sessionId, tool.tagNumber, "full"); + updateSessionMeta(db, sessionId, { + clearedReasoningThroughTag: Math.max(...tags.map((tag) => tag.tagNumber)), + }); + expect(getPendingOps(db, sessionId)).toEqual([]); + }; + return { + db, + pass, + restart, + seedLegacy, + close: () => { + clearContextHandlerSession(sessionId); + db.close(); + }, + }; +} + +describe("native upgrade application", () => { + it("keeps legacy native bytes across defer passes, then persists one authorized transition", async () => { + const sessionId = "ses-native-upgrade"; + const f = fixture(sessionId); + try { + await f.seedLegacy(); + const first = await f.pass(0); + const a = nativeBytes(first); + expect(a).toContain(staleInput); + expect(a).toContain(ciphertext); + expect(first[1]).toMatchObject({ + content: [ + { type: "thinking", thinking: "" }, + {}, + { arguments: { dropped: expect.any(String) } }, + ], + }); + expect(nativeBytes(await f.pass(0))).toBe(a); + f.restart(); + expect(nativeBytes(await f.pass(0))).toBe(a); + expect(getNativeToolInputs(f.db, sessionId).size).toBe(0); + expect(getNativeReasoningIds(f.db, sessionId).size).toBe(0); + + const b = nativeBytes(await f.pass(90)); + expect(b).not.toContain(staleInput); + expect(b).not.toContain(ciphertext); + expect(b).toContain("visible answer"); + expect(b).toContain("image_generation_call"); + expect(b).toContain("call-old"); + expect(getNativeToolInputs(f.db, sessionId).has(callId)).toBe(true); + expect(getNativeReasoningIds(f.db, sessionId).has("entry-old")).toBe( + true, + ); + expect(nativeBytes(await f.pass(0))).toBe(b); + f.restart(); + expect(nativeBytes(await f.pass(0))).toBe(b); + } finally { + f.close(); + } + }); + + for (const column of [ + "pi_native_tool_inputs", + "pi_native_reasoning_ids", + ] as const) { + it(`does not publish failed ${column} activation or retry it on defer`, async () => { + const sessionId = `ses-native-failure-${column}`; + const f = fixture(sessionId); + try { + await f.seedLegacy(); + f.db.exec( + `CREATE TRIGGER fail_native_write BEFORE UPDATE OF ${column} ON session_meta BEGIN SELECT RAISE(FAIL, 'native persistence failure'); END`, + ); + const failed = nativeBytes(await f.pass(90)); + const retained = + column === "pi_native_tool_inputs" ? staleInput : ciphertext; + expect(failed).toContain(retained); + f.db.exec("DROP TRIGGER fail_native_write"); + expect(nativeBytes(await f.pass(0))).toBe(failed); + f.restart(); + expect(nativeBytes(await f.pass(0))).toBe(failed); + const retried = nativeBytes(await f.pass(90)); + expect(retried).not.toContain(retained); + expect(nativeBytes(await f.pass(0))).toBe(retried); + } finally { + f.close(); + } + }); + } + + it("freezes exact tool input values until another authorized mutation", () => { + const db = createTestDb(); + const sessionId = "ses-native-input-progress"; + try { + const run = (marker: string, canApply: boolean, changed = true) => { + const message = oldAssistant() as unknown as { + content: Array<{ type: string; arguments?: Record }>; + }; + message.content[2].arguments = { dropped: marker }; + const messages: unknown[] = [message]; + applyNativeToolInputReplayPi({ + db, + sessionId, + messages, + canApply, + changes: changed ? new Map([[0, new Set([callId])]]) : new Map(), + }); + return nativeBytes(messages); + }; + const b = run("first marker", true); + expect(b).toContain("first marker"); + expect(run("different marker", false)).toBe(b); + expect(run("different marker", false, false)).toBe(b); + db.exec( + "CREATE TRIGGER fail_update BEFORE UPDATE OF pi_native_tool_inputs ON session_meta BEGIN SELECT RAISE(FAIL, 'rejected native update'); END", + ); + expect(run("different marker", true)).toBe(b); + db.exec("DROP TRIGGER fail_update"); + expect(run("different marker", false)).toBe(b); + const c = run("different marker", true); + expect(c).toContain("different marker"); + expect(run("first marker", false)).toBe(c); + } finally { + db.close(); + } + }); + + it("waits for a real entry identity before persisting native reasoning", () => { + const db = createTestDb(); + const sessionId = "ses-native-unresolved-id"; + try { + const run = (id: string, canApply: boolean) => { + const messages: unknown[] = [oldAssistant()]; + applyNativeReasoningReplayPi({ + db, + sessionId, + messages, + stableId: () => id, + messageIdToMaxTag: new Map([[id, 1]]), + localWatermark: 10, + clearReasoningAge: 100, + omissionAllowed: true, + detectAged: false, + canApply, + }); + return nativeBytes(messages); + }; + const unresolved = run("pi-msg-0-2-assistant", true); + expect(unresolved).toContain(ciphertext); + expect(getNativeReasoningIds(db, sessionId).size).toBe(0); + expect(run("entry-old", false)).toBe(unresolved); + expect(run("entry-old", false)).toBe(unresolved); + const cleared = run("entry-old", true); + expect(cleared).not.toContain(ciphertext); + expect(run("entry-old", false)).toBe(cleared); + } finally { + db.close(); + } + }); + + it("persists native-only reasoning without borrowing or advancing the local watermark", () => { + const db = createTestDb(); + const sessionId = "ses-native-only"; + try { + const original = oldAssistant() as unknown as Record; + original.content = [{ type: "text", text: "visible answer" }]; + const before = nativeBytes([original]); + const run = (canApply: boolean) => { + const messages = [original]; + applyNativeReasoningReplayPi({ + db, + sessionId, + messages, + messageIdToMaxTag: new Map([ + ["entry-old", 1], + ["entry-new", 100], + ]), + stableId: () => "entry-old", + localWatermark: 0, + clearReasoningAge: 10, + omissionAllowed: true, + detectAged: canApply, + canApply, + }); + return nativeBytes(messages); + }; + const after = run(true); + expect(after).not.toContain(ciphertext); + expect(nativeBytes([original])).toBe(before); + expect( + getOrCreateSessionMeta(db, sessionId).clearedReasoningThroughTag, + ).toBe(0); + expect(run(false)).toBe(after); + expect(run(false)).toBe(after); + } finally { + db.close(); + } + }); + + it("does not let tool activation authorize reasoning or age into new reasoning on defer", () => { + const db = createTestDb(); + const sessionId = "ses-native-independent"; + try { + getOrCreateSessionMeta(db, sessionId); + const toolMessage = oldAssistant() as unknown as { + content: Array<{ arguments?: Record }>; + }; + toolMessage.content[2].arguments = { dropped: "native tool marker" }; + applyNativeToolInputReplayPi({ + db, + sessionId, + messages: [toolMessage], + changes: new Map([[0, new Set([callId])]]), + canApply: true, + }); + expect(getNativeToolInputs(db, sessionId).has(callId)).toBe(true); + const run = ( + omissionAllowed: boolean, + canApply: boolean, + localWatermark: number, + ) => { + const messages: unknown[] = [oldAssistant()]; + applyNativeReasoningReplayPi({ + db, + sessionId, + messages, + messageIdToMaxTag: new Map([ + ["old", 1], + ["new", 200], + ]), + stableId: () => "old", + localWatermark, + clearReasoningAge: 10, + omissionAllowed, + canApply, + detectAged: false, + }); + return nativeBytes(messages); + }; + expect(run(false, true, 10)).toContain(ciphertext); + expect(getNativeReasoningIds(db, sessionId).size).toBe(0); + expect(run(true, false, 10)).toContain(ciphertext); + expect(run(true, true, 0)).toContain(ciphertext); + const cleared = run(true, true, 10); + expect(cleared).not.toContain(ciphertext); + expect(run(true, false, 0)).toBe(cleared); + } finally { + db.close(); + } + }); +}); diff --git a/packages/pi-plugin/src/native-replay-state-pi.ts b/packages/pi-plugin/src/native-replay-state-pi.ts new file mode 100644 index 000000000..c2ca900fe --- /dev/null +++ b/packages/pi-plugin/src/native-replay-state-pi.ts @@ -0,0 +1,157 @@ +import type { ContextDatabase } from "@magic-context/core/features/magic-context/storage"; +import { + addNativeReasoningIds, + getNativeReasoningIds, + getNativeToolInputs, + saveNativeToolInputs, +} from "@magic-context/core/features/magic-context/storage-native-replay"; +import { sessionLog } from "@magic-context/core/shared/logger"; +import { isRecord } from "@magic-context/core/shared/record-type-guard"; +import { + clearNativeReasoning, + rewriteNativeToolInput, +} from "./native-replay-pi"; +import { SYNTH_USER_ID_PREFIX } from "./read-session-pi"; + +type ReplayArgs = { + db: ContextDatabase; + sessionId: string; + messages: unknown[]; +}; + +/** Replay frozen inputs first; publish new native bytes only after their write succeeds. */ +export function applyNativeToolInputReplayPi( + args: ReplayArgs & { + changes: ReadonlyMap>; + canApply: boolean; + }, +): number { + const saved = getNativeToolInputs(args.db, args.sessionId); + if (saved.size === 0 && (!args.canApply || args.changes.size === 0)) return 0; + const nextInputs = new Map(); + const pending = new Map>(); + + for (let index = 0; index < args.messages.length; index++) { + const original = args.messages[index]; + if ( + !isRecord(original) || + original.role !== "assistant" || + !Array.isArray(original.content) || + original.providerPayload == null + ) + continue; + const calls = original.content.filter( + (part): part is Record => + isRecord(part) && + part.type === "toolCall" && + typeof part.id === "string", + ); + let replay = original; + for (const call of calls) { + const input = saved.get(call.id as string); + if (input === undefined) continue; + const parsed: unknown = JSON.parse(input); + if (!isRecord(parsed)) + throw new Error("Invalid persisted native tool input"); + const candidate = { ...replay }; + rewriteNativeToolInput(candidate, call.id as string, parsed); + if (candidate.providerPayload !== replay.providerPayload) + replay = candidate; + } + args.messages[index] = replay; + if (!args.canApply) continue; + const changed = args.changes.get(index); + if (!changed) continue; + + let next = replay; + for (const call of calls) { + const id = call.id as string; + if (!changed.has(id) || !isRecord(call.arguments)) continue; + let serialized: string | undefined; + try { + serialized = JSON.stringify(call.arguments); + } catch { + continue; + } + if (serialized === undefined || saved.get(id) === serialized) continue; + const normalized: unknown = JSON.parse(serialized); + if (!isRecord(normalized)) continue; + const candidate = { ...next }; + rewriteNativeToolInput(candidate, id, normalized); + if (candidate.providerPayload === next.providerPayload) continue; + next = candidate; + nextInputs.set(id, serialized); + } + if (next !== replay) pending.set(index, next); + } + + if (nextInputs.size === 0) return 0; + try { + saveNativeToolInputs(args.db, args.sessionId, nextInputs); + } catch (error) { + sessionLog( + args.sessionId, + `native input activation failed; retaining previous replay: ${error instanceof Error ? error.message : String(error)}`, + ); + return 0; + } + for (const [index, message] of pending) args.messages[index] = message; + return nextInputs.size; +} + +/** Native reasoning has its own durable decisions, independent of the local watermark. */ +export function applyNativeReasoningReplayPi( + args: ReplayArgs & { + messageIdToMaxTag: ReadonlyMap; + stableId: (message: unknown, index: number) => string | undefined; + localWatermark: number; + clearReasoningAge: number; + omissionAllowed: boolean; + canApply: boolean; + detectAged: boolean; + }, +): number { + if (!args.omissionAllowed) return 0; + const saved = getNativeReasoningIds(args.db, args.sessionId); + let maxTag = 0; + for (const tag of args.messageIdToMaxTag.values()) + maxTag = Math.max(maxTag, tag); + const cutoff = args.detectAged + ? Math.max(args.localWatermark, maxTag - args.clearReasoningAge) + : args.localWatermark; + const nextIds = new Set(); + const pending = new Map>(); + + for (let index = 0; index < args.messages.length; index++) { + const original = args.messages[index]; + if (!isRecord(original) || original.role !== "assistant") continue; + const id = args.stableId(original, index); + // Positional fallback IDs drift before adoption; they cannot own durable replay. + if (!id || id.startsWith("pi-msg-") || id.startsWith(SYNTH_USER_ID_PREFIX)) + continue; + const replay = saved.has(id); + const tag = args.messageIdToMaxTag.get(id) ?? 0; + if (!replay && (!args.canApply || tag === 0 || tag > cutoff)) continue; + const candidate = { ...original }; + if (clearNativeReasoning(candidate, true) !== "cleared") continue; + if (replay) { + args.messages[index] = candidate; + } else { + nextIds.add(id); + pending.set(index, candidate); + } + } + + if (nextIds.size === 0) return 0; + try { + addNativeReasoningIds(args.db, args.sessionId, nextIds); + } catch (error) { + sessionLog( + args.sessionId, + `native reasoning activation failed; retaining previous replay: ${error instanceof Error ? error.message : String(error)}`, + ); + return 0; + } + for (const [index, message] of pending) args.messages[index] = message; + return nextIds.size; +} diff --git a/packages/pi-plugin/src/reasoning-replay-pi.test.ts b/packages/pi-plugin/src/reasoning-replay-pi.test.ts index 344c80126..067ac22ce 100644 --- a/packages/pi-plugin/src/reasoning-replay-pi.test.ts +++ b/packages/pi-plugin/src/reasoning-replay-pi.test.ts @@ -128,71 +128,52 @@ describe("clearOldReasoningPi", () => { }); }); - it("clears native-only reasoning and replays the same durable watermark", () => { - const db = makeDb(); - const sessionId = "ses-native-reasoning"; - try { - const native = { - type: "openaiResponsesHistory", - dt: true, - items: [ - { type: "reasoning", encrypted_content: "persisted-only-reasoning" }, - { - type: "function_call", - id: "fc1", - call_id: "call1", - name: "read", - arguments: "{}", - }, - ], - }; - const original = { - role: "assistant", - timestamp: 1, - content: [ - { type: "thinking", thinking: "" }, - { type: "toolCall", id: "call1|fc1", name: "read", arguments: {} }, - ], - providerPayload: native, - }; - const first = structuredClone(original); - const messageIdToMaxTag = new Map(); - messageIdToMaxTag.set("a", 1); - messageIdToMaxTag.set("recent", 10); - const options = { - messageIdToMaxTag, - piMessageStableId: () => "a", - nativeReasoningMayClear: true, - }; - const cleared = clearOldReasoningPi({ - ...options, - messages: [first], - clearReasoningAge: 3, - }); - expect(cleared).toEqual({ cleared: 1, newWatermark: 1 }); - expect(first.providerPayload.items).toEqual([native.items[1]]); + it("leaves native-only reasoning to the native coordinator without advancing the local watermark", () => { + const native = { + type: "openaiResponsesHistory", + dt: true, + items: [ + { type: "reasoning", encrypted_content: "persisted-only-reasoning" }, + { + type: "function_call", + id: "fc1", + call_id: "call1", + name: "read", + arguments: "{}", + }, + ], + }; + const nativeSnapshot = structuredClone(native); + const localMessage = { + role: "assistant", + timestamp: 1, + content: [{ type: "thinking", thinking: "local thinking" }], + }; + const nativeOnlyMessage = { + role: "assistant", + timestamp: 2, + content: [ + { type: "toolCall", id: "call1|fc1", name: "read", arguments: {} }, + ], + providerPayload: native, + }; + const messages = [localMessage, nativeOnlyMessage]; + const result = clearOldReasoningPi({ + messages, + messageIdToMaxTag: new Map([ + ["local", 1], + ["native", 2], + ["recent", 10], + ]), + clearReasoningAge: 3, + piMessageStableId: (_message, index) => + index === 0 ? "local" : "native", + }); - getOrCreateSessionMeta(db, sessionId); - updateSessionMeta(db, sessionId, { - clearedReasoningThroughTag: cleared.newWatermark, - }); - const resumed = structuredClone(original); - expect( - replayClearedReasoningPi({ - ...options, - messages: [resumed], - db, - sessionId, - }), - ).toBe(1); - expect(resumed.providerPayload).toEqual(first.providerPayload); - expect(native.items[0]).toEqual({ - type: "reasoning", - encrypted_content: "persisted-only-reasoning", - }); - } finally { - db.close(); - } + expect(result).toEqual({ cleared: 1, newWatermark: 1 }); + expect(localMessage.content).toEqual([{ type: "thinking", thinking: "" }]); + expect(nativeOnlyMessage.providerPayload).toBe(native); + expect(nativeOnlyMessage.providerPayload).toEqual(nativeSnapshot); }); it("preserves a full native snapshot while clearing local Pi thinking and replaying its watermark", () => { @@ -223,7 +204,6 @@ describe("clearOldReasoningPi", () => { ]); const options = { messageIdToMaxTag, - nativeReasoningMayClear: true, piMessageStableId: () => "a", }; @@ -260,103 +240,6 @@ describe("clearOldReasoningPi", () => { } }); - it("cleans local thinking when native clearing is disallowed or native reasoning has plaintext", () => { - const db = makeDb(); - try { - for (const fixture of [ - { - label: "native-clearing-disallowed", - nativeReasoningMayClear: false, - providerPayload: { - type: "openaiResponsesHistory", - dt: true, - items: [ - { - type: "reasoning", - encrypted_content: "opaque-encrypted-content", - }, - ], - }, - }, - { - label: "plaintext-native-reasoning", - nativeReasoningMayClear: true, - providerPayload: { - type: "openaiResponsesHistory", - dt: true, - items: [ - { - type: "reasoning", - encrypted_content: "opaque-encrypted-content", - content: [ - { - type: "reasoning_text", - text: "required plaintext", - }, - ], - }, - ], - }, - }, - ] as const) { - const original = { - role: "assistant", - content: [ - { - type: "thinking", - thinking: "ordinary local thinking", - thinkingSignature: "local-signature", - }, - { type: "text", text: "reply" }, - ], - providerPayload: fixture.providerPayload, - }; - const nativeSnapshot = structuredClone(fixture.providerPayload); - const messageIdToMaxTag = new Map([ - ["a", 1], - ["recent", 4], - ]); - const options = { - messageIdToMaxTag, - nativeReasoningMayClear: fixture.nativeReasoningMayClear, - piMessageStableId: () => "a", - }; - - const first = structuredClone(original); - const cleared = clearOldReasoningPi({ - ...options, - messages: [first], - clearReasoningAge: 2, - }); - expect(cleared).toEqual({ cleared: 1, newWatermark: 1 }); - expect(first.content).toEqual([ - { type: "thinking", thinking: "" }, - { type: "text", text: "reply" }, - ]); - expect(first.providerPayload).toEqual(nativeSnapshot); - - const sessionId = `ses-${fixture.label}`; - getOrCreateSessionMeta(db, sessionId); - updateSessionMeta(db, sessionId, { - clearedReasoningThroughTag: cleared.newWatermark, - }); - const resumed = structuredClone(original); - expect( - replayClearedReasoningPi({ - ...options, - messages: [resumed], - db, - sessionId, - }), - ).toBe(1); - expect(resumed).toEqual(first); - expect(resumed.providerPayload).toEqual(nativeSnapshot); - } - } finally { - db.close(); - } - }); - it("does nothing when ageCutoff is 0 or below", () => { const messages = [ { @@ -417,7 +300,6 @@ describe("clearOldReasoningPi", () => { const options = { messageIdToMaxTag, piMessageStableId, - nativeReasoningMayClear: true, }; const first = structuredClone(original); diff --git a/packages/pi-plugin/src/reasoning-replay-pi.ts b/packages/pi-plugin/src/reasoning-replay-pi.ts index c65ec4f7f..a042d17c5 100644 --- a/packages/pi-plugin/src/reasoning-replay-pi.ts +++ b/packages/pi-plugin/src/reasoning-replay-pi.ts @@ -37,7 +37,6 @@ import type { ContextDatabase } from "@magic-context/core/features/magic-context/storage"; import { getOrCreateSessionMeta } from "@magic-context/core/features/magic-context/storage"; import type { TagTarget } from "@magic-context/core/hooks/magic-context/tag-messages"; -import { clearNativeReasoning } from "./native-replay-pi"; type PiTextContent = { type: "text"; text: string }; type PiThinkingContent = { @@ -110,10 +109,10 @@ export function buildMessageIdToMaxTag( } /** - * Clear typed reasoning on assistant messages whose tag number is - * older than `(maxTag - clearReasoningAge)`. Returns the highest tag - * number that was actually cleared, so the caller can persist the - * watermark via `setReasoningWatermark`. + * Clear local typed reasoning on assistant messages whose tag number is older + * than `(maxTag - clearReasoningAge)`. Returns the highest tag number that was + * actually cleared, so the caller can persist the local watermark via + * `setReasoningWatermark`. * * Mirrors OpenCode's `clearOldReasoning` (strip-content.ts). */ @@ -121,7 +120,6 @@ export function clearOldReasoningPi(args: { messages: unknown[]; messageIdToMaxTag: Map; clearReasoningAge: number; - nativeReasoningMayClear?: boolean; piMessageStableId: (msg: unknown, index: number) => string | undefined; }): { cleared: number; newWatermark: number } { const { messages, messageIdToMaxTag, clearReasoningAge, piMessageStableId } = @@ -148,14 +146,7 @@ export function clearOldReasoningPi(args: { const msgTag = messageIdToMaxTag.get(id) ?? 0; if (msgTag === 0 || msgTag > ageCutoff) continue; - const nativeReasoning = clearNativeReasoning( - msg, - args.nativeReasoningMayClear === true, - ); - // Native preservation protects providerPayload only; local Pi thinking - // still follows the ordinary per-part cleanup policy. - const clearedBefore = cleared; - + let clearedThisMessage = false; for (const part of msg.content) { if ( part && @@ -180,12 +171,13 @@ export function clearOldReasoningPi(args: { tp.thinking = CLEARED; tp.thinkingSignature = undefined; cleared++; + clearedThisMessage = true; } } } - if (nativeReasoning === "cleared" && cleared === clearedBefore) cleared++; - - if (cleared > 0 && msgTag > newWatermark) newWatermark = msgTag; + if (clearedThisMessage && msgTag > newWatermark) { + newWatermark = msgTag; + } } return { cleared, newWatermark }; @@ -252,17 +244,15 @@ export function stripInlineThinkingPi(args: { } /** - * Replay typed-reasoning clearing on EVERY pass (execute or defer). - * Mirrors OpenCode's `replayClearedReasoning` — required for cache - * stability so the Pi assistant content array stays byte-identical - * across passes. + * Replay local typed-reasoning clearing on EVERY pass (execute or defer). + * Mirrors OpenCode's `replayClearedReasoning` — required for cache stability + * so the Pi assistant content array stays byte-identical across passes. */ export function replayClearedReasoningPi(args: { db: ContextDatabase; sessionId: string; messages: unknown[]; messageIdToMaxTag: Map; - nativeReasoningMayClear?: boolean; piMessageStableId: (msg: unknown, index: number) => string | undefined; }): number { const { db, sessionId, messages, messageIdToMaxTag, piMessageStableId } = @@ -284,14 +274,6 @@ export function replayClearedReasoningPi(args: { const msgTag = messageIdToMaxTag.get(id) ?? 0; if (msgTag === 0 || msgTag > watermark) continue; - const nativeReasoning = clearNativeReasoning( - msg, - args.nativeReasoningMayClear === true, - ); - // Native preservation protects providerPayload only; local Pi thinking - // still follows the ordinary per-part cleanup policy. - const clearedBefore = cleared; - for (const part of msg.content) { if ( part && @@ -313,7 +295,6 @@ export function replayClearedReasoningPi(args: { } } } - if (nativeReasoning === "cleared" && cleared === clearedBefore) cleared++; } return cleared; } diff --git a/packages/pi-plugin/src/transcript-pi.test.ts b/packages/pi-plugin/src/transcript-pi.test.ts index 524cf1547..621f81762 100644 --- a/packages/pi-plugin/src/transcript-pi.test.ts +++ b/packages/pi-plugin/src/transcript-pi.test.ts @@ -75,89 +75,136 @@ describe("createPiTranscript", () => { expect(transcript.getOutputMessages()).toBe(messages); }); - it("keeps native text stable while reducing a paired tool input", () => { - const db = createTestDb(); - try { - const text = " Keep the indentation and the explanation"; - const reasoning = { - type: "reasoning", - encrypted_content: "unrelated-encrypted-reasoning", - }; - const hostedOutput = { - type: "image_generation_call", - result: "image-data", - }; - const native = { - type: "openaiResponsesHistory", - dt: true, - items: [ - reasoning, - hostedOutput, + it("defers native tool input replay while retaining current Pi mutations", () => { + const text = "Keep the native response untouched until authorization"; + const reasoning = { + type: "reasoning", + encrypted_content: "unrelated-encrypted-reasoning", + }; + const hostedOutput = { + type: "image_generation_call", + result: "image-data", + }; + const native = { + type: "openaiResponsesHistory", + dt: true, + items: [ + reasoning, + hostedOutput, + { + type: "message", + role: "assistant", + id: "msg-answer", + content: [{ type: "output_text", text }], + }, + { + type: "function_call", + id: "fc-sentinel", + call_id: "call-sentinel", + name: "write", + arguments: JSON.stringify({ content: "large original input" }), + }, + { + type: "function_call", + id: "fc-text", + call_id: "call-text", + name: "edit", + arguments: JSON.stringify({ patch: "large original patch" }), + }, + { + type: "function_call", + id: "fc-input", + call_id: "call-input", + name: "read", + arguments: JSON.stringify({ path: "large original file" }), + }, + ], + }; + const nativeSnapshot = structuredClone(native); + const original = { + ...assistantMessage(text, 1, { + content: [ + { type: "text", text, textSignature: "msg-answer" }, { - type: "message", - role: "assistant", - id: "msg-answer", - content: [{ type: "output_text", text }], + type: "toolCall", + id: "call-sentinel|fc-sentinel", + name: "write", + arguments: { content: "large original input" }, }, { - type: "function_call", - id: "fc-write", - call_id: "call-write", - name: "write", - arguments: JSON.stringify({ content: "large original input" }), + type: "toolCall", + id: "call-text|fc-text", + name: "edit", + arguments: { patch: "large original patch" }, + }, + { + type: "toolCall", + id: "call-input|fc-input", + name: "read", + arguments: { path: "large original file" }, }, ], - }; - const original = { - ...assistantMessage(text, 1, { - content: [ - { type: "text", text, textSignature: "msg-answer" }, - { - type: "toolCall", - id: "call-write|fc-write", - name: "write", - arguments: { content: "large original input" }, - }, - ], - }), - providerPayload: native, - }; - const messages = [ - original, - toolResultMessage("call-write|fc-write", "written", 2), - ]; - const sessionId = "ses-native-replay"; - const tagger = createTagger(); - tagger.initFromDb(sessionId, db); - const tagged = createPiTranscript(messages, sessionId); - tagTranscript(sessionId, tagged, tagger, db); - tagged.commit(); - const taggedAssistant = messages[0] as typeof original; - expect(textOf(taggedAssistant)).not.toBe(text); - expect(taggedAssistant.providerPayload).toBe(native); - - const reduced = createPiTranscript(messages, sessionId); - const toolPart = reduced.messages[0]?.parts[1]; - if (!toolPart) throw new Error("missing tool part"); - toolPart.replaceWithSentinel("[dropped input]"); - reduced.commit(); - const payload = (messages[0] as typeof original).providerPayload; - expect(payload.items[2]).toEqual(native.items[2]); - expect(payload.items[3]).toMatchObject({ - type: "function_call", - id: "fc-write", - call_id: "call-write", + }), + providerPayload: native, + }; + const messages = [userMessage("prior context", 0), original]; + const assistantMessageIndex = 1; + const transcript = createPiTranscript(messages, "ses-native-replay"); + const parts = transcript.messages[assistantMessageIndex]?.parts ?? []; + + expect(parts[0]?.setText("ordinary content changed")).toBe(true); + expect(parts[1]?.replaceWithSentinel("[dropped input]")).toBe(true); + expect(parts[2]?.setText("[truncated input]")).toBe(true); + expect(parts[3]?.setToolInput?.({ path: "short.txt" })).toBe(true); + expect( + Array.from( + transcript.getToolInputChanges().get(assistantMessageIndex) ?? [], + ), + ).toEqual([ + "call-sentinel|fc-sentinel", + "call-text|fc-text", + "call-input|fc-input", + ]); + + transcript.commit(); + const output = transcript.getOutputMessages() as Array<{ + content: Array<{ + type: string; + text?: string; + id?: string; + arguments?: Record; + }>; + providerPayload: typeof native; + }>; + expect(output[assistantMessageIndex]?.content).toEqual([ + { + type: "text", + text: "ordinary content changed", + textSignature: "msg-answer", + }, + { + type: "toolCall", + id: "call-sentinel|fc-sentinel", name: "write", - arguments: JSON.stringify({ dropped: "[dropped input]" }), - }); - expect(payload.items[0]).toEqual(reasoning); - expect(payload.items[1]).toEqual(hostedOutput); - expect(native.items[3]).toMatchObject({ - arguments: JSON.stringify({ content: "large original input" }), - }); - } finally { - closeQuietly(db); - } + arguments: { dropped: "[dropped input]" }, + }, + { + type: "toolCall", + id: "call-text|fc-text", + name: "edit", + arguments: { __magic_context_replacement__: "[truncated input]" }, + }, + { + type: "toolCall", + id: "call-input|fc-input", + name: "read", + arguments: { path: "short.txt" }, + }, + ]); + expect(output[assistantMessageIndex]?.providerPayload).toBe(native); + expect(output[assistantMessageIndex]?.providerPayload).toEqual( + nativeSnapshot, + ); }); it("leaves Pi and OMP whitespace-only assistant framing untagged after peeling MC tags", () => { diff --git a/packages/pi-plugin/src/transcript-pi.ts b/packages/pi-plugin/src/transcript-pi.ts index edec56fa5..c2f43a666 100644 --- a/packages/pi-plugin/src/transcript-pi.ts +++ b/packages/pi-plugin/src/transcript-pi.ts @@ -77,7 +77,6 @@ import type { TranscriptPart, TranscriptPartKind, } from "@magic-context/core/shared/transcript"; -import { rewriteNativeToolInput } from "./native-replay-pi"; import { resolvePiHarnessKind } from "./pi-harness-kind"; import { resolvePiStableId, SYNTH_USER_ID_PREFIX } from "./read-session-pi"; @@ -133,6 +132,7 @@ type PiToolResultMessage = { }; type PiAgentMessage = PiUserMessage | PiAssistantMessage | PiToolResultMessage; +type MarkDirty = (messageIndex: number, toolCallId?: string) => void; /** * Wrap a Pi `AgentMessage[]` as a Transcript. Builds the normalized @@ -164,8 +164,8 @@ export function createPiTranscript( * (tagging, drops, caveman) write to and `commit()` flushes back to source. * * Phases that mutate messages OUTSIDE the transcript part API — reasoning - * clearing/replay, which set `part.thinking = "[cleared]"` in place — MUST - * target this array, not the original `source`. Tagging/drops/caveman + * clearing/replay, which empty local `part.thinking` in place — MUST target + * this array, not the original `source`. Tagging/drops/caveman * REASSIGN `working[idx]` to fresh spread-copied objects; if reasoning mutated * `source[idx]` (a now-divergent object) instead, the later `commit()` would * overwrite `source[idx] = working[idx]` and silently discard the reasoning @@ -175,9 +175,16 @@ export function createPiTranscript( * mutation in the single channel `commit()` flushes. */ getWorkingMessages(): PiAgentMessage[]; + /** + * Pi-only change inventory for the native replay stage. Keys are current + * working-array indices; values are tool-call IDs whose arguments changed + * through a transcript part setter. + */ + getToolInputChanges(): ReadonlyMap>; } { const working = source.slice() as unknown as PiAgentMessage[]; const dirtyMessages = new Set(); + const toolInputChanges = new Map>(); // Normalize: fold consecutive toolResult runs into the immediately // following user message as tool_result transcript parts. Track @@ -185,8 +192,15 @@ export function createPiTranscript( const transcriptMessages: TranscriptMessage[] = buildTranscriptView( working, sessionId, - (messageIndex) => { + (messageIndex, toolCallId) => { dirtyMessages.add(messageIndex); + if (toolCallId === undefined) return; + let toolCallIds = toolInputChanges.get(messageIndex); + if (toolCallIds === undefined) { + toolCallIds = new Set(); + toolInputChanges.set(messageIndex, toolCallIds); + } + toolCallIds.add(toolCallId); }, entryIds, ); @@ -228,6 +242,9 @@ export function createPiTranscript( getWorkingMessages(): PiAgentMessage[] { return working; }, + getToolInputChanges(): ReadonlyMap> { + return toolInputChanges; + }, }; } @@ -241,7 +258,7 @@ export function createPiTranscript( function buildTranscriptView( working: PiAgentMessage[], sessionId: string | undefined, - markDirty: (messageIndex: number) => void, + markDirty: MarkDirty, entryIds: readonly (string | undefined)[] | undefined, ): TranscriptMessage[] { const result: TranscriptMessage[] = []; @@ -343,7 +360,7 @@ function createUserTranscriptMessage( index: number, sessionId: string | undefined, foldedToolResults: { msg: PiToolResultMessage; index: number }[], - markDirty: (messageIndex: number) => void, + markDirty: MarkDirty, entryIds: readonly (string | undefined)[] | undefined, ): TranscriptMessage { const userMsg = working[index] as PiUserMessage; @@ -405,7 +422,7 @@ function createSyntheticToolResultUserMessage( working: PiAgentMessage[], sessionId: string | undefined, toolResultRun: { msg: PiToolResultMessage; index: number }[], - markDirty: (messageIndex: number) => void, + markDirty: MarkDirty, entryIds: readonly (string | undefined)[] | undefined, ): TranscriptMessage { return { @@ -450,7 +467,7 @@ function createAssistantTranscriptMessage( working: PiAgentMessage[], index: number, sessionId: string | undefined, - markDirty: (messageIndex: number) => void, + markDirty: MarkDirty, entryIds: readonly (string | undefined)[] | undefined, ): TranscriptMessage { const msg = working[index] as PiAssistantMessage; @@ -472,7 +489,7 @@ function createOpaqueTranscriptMessage( working: PiAgentMessage[], index: number, sessionId: string | undefined, - _markDirty: (messageIndex: number) => void, + _markDirty: MarkDirty, entryIds: readonly (string | undefined)[] | undefined, ): TranscriptMessage { const msg = working[index]; @@ -500,7 +517,7 @@ function createOpaqueTranscriptMessage( function createPiUserStringPart( working: PiAgentMessage[], messageIndex: number, - markDirty: (messageIndex: number) => void, + markDirty: MarkDirty, ): TranscriptPart { return { kind: "text", @@ -544,7 +561,7 @@ function createPiUserArrayPart( working: PiAgentMessage[], messageIndex: number, partIndex: number, - markDirty: (messageIndex: number) => void, + markDirty: MarkDirty, ): TranscriptPart { const msg = working[messageIndex] as PiUserMessage; const part = Array.isArray(msg.content) ? msg.content[partIndex] : undefined; @@ -605,7 +622,7 @@ function createPiAssistantPart( working: PiAgentMessage[], messageIndex: number, partIndex: number, - markDirty: (messageIndex: number) => void, + markDirty: MarkDirty, ): TranscriptPart { const msg = working[messageIndex] as PiAssistantMessage; const part = msg.content[partIndex]; @@ -675,8 +692,7 @@ function createPiAssistantPart( ...(working[messageIndex] as PiAssistantMessage), content: newContent, }; - rewriteNativeToolInput(working[messageIndex], p.id, replacementArgs); - markDirty(messageIndex); + markDirty(messageIndex, p.id); return true; } return false; @@ -735,8 +751,7 @@ function createPiAssistantPart( ...(working[messageIndex] as PiAssistantMessage), content: newContent, }; - rewriteNativeToolInput(working[messageIndex], p.id, input); - markDirty(messageIndex); + markDirty(messageIndex, p.id); return true; }, // Replace this assistant part's content with a sentinel placeholder. @@ -780,15 +795,10 @@ function createPiAssistantPart( ...(working[messageIndex] as PiAssistantMessage), content: newContent, }; - const replacement = newContent[partIndex]; - if (existing?.type === "toolCall" && replacement?.type === "toolCall") { - rewriteNativeToolInput( - working[messageIndex], - existing.id, - replacement.arguments, - ); - } - markDirty(messageIndex); + markDirty( + messageIndex, + existing?.type === "toolCall" ? existing.id : undefined, + ); return true; }, }; @@ -798,7 +808,7 @@ function createPiToolResultPart( working: PiAgentMessage[], messageIndex: number, partIndex: number, - markDirty: (messageIndex: number) => void, + markDirty: MarkDirty, ): TranscriptPart { const msg = working[messageIndex] as PiToolResultMessage; // Every block inside one Pi ToolResultMessage belongs to the same diff --git a/packages/plugin/scripts/clone-session.test.ts b/packages/plugin/scripts/clone-session.test.ts index da89dbd9c..d3f81ef0b 100644 --- a/packages/plugin/scripts/clone-session.test.ts +++ b/packages/plugin/scripts/clone-session.test.ts @@ -489,6 +489,45 @@ describe("clone-session", () => { mutable.close(); }); + it("keeps filtered native decisions after the generic metadata copy", () => { + const fixture = makeFixture(); + const frozenInput = '{"path":"msg_source_1","dropped":"marker"}'; + const source = new Database(fixture.contextPath); + try { + source.exec(` + ALTER TABLE session_meta ADD COLUMN pi_native_tool_inputs TEXT; + ALTER TABLE session_meta ADD COLUMN pi_native_reasoning_ids TEXT; + `); + source.prepare( + "UPDATE session_meta SET pi_native_tool_inputs = ?, pi_native_reasoning_ids = ? WHERE session_id = ?", + ).run( + JSON.stringify({ tool_source_1: frozenInput, outside_call: '{"dropped":"outside"}' }), + JSON.stringify(["msg_source_2", "outside_assistant"]), + fixture.sourceSessionId, + ); + } finally { + source.close(); + } + const result = cloneSession({ + sessionId: fixture.sourceSessionId, + opencodeDbPath: fixture.opencodePath, + contextDbPath: fixture.contextPath, + }); + const clone = new Database(fixture.contextPath, { readonly: true }); + try { + const tool = clone.prepare( + "SELECT message_id, tool_owner_message_id FROM tags WHERE session_id = ? AND type = 'tool'", + ).get(result.plan.destinationSessionId) as { message_id: string; tool_owner_message_id: string }; + const state = clone.prepare( + "SELECT pi_native_tool_inputs, pi_native_reasoning_ids FROM session_meta WHERE session_id = ?", + ).get(result.plan.destinationSessionId) as { pi_native_tool_inputs: string; pi_native_reasoning_ids: string }; + expect(JSON.parse(state.pi_native_tool_inputs)).toEqual({ [tool.message_id]: frozenInput }); + expect(JSON.parse(state.pi_native_reasoning_ids)).toEqual([tool.tool_owner_message_id]); + } finally { + clone.close(); + } + }); + it("clones notes before project authority is attached and leaves the guard active", () => { const fixture = makeFixture(); const context = new Database(fixture.contextPath); diff --git a/packages/plugin/scripts/clone-session.ts b/packages/plugin/scripts/clone-session.ts index c8a1873e4..978efb730 100644 --- a/packages/plugin/scripts/clone-session.ts +++ b/packages/plugin/scripts/clone-session.ts @@ -1055,8 +1055,13 @@ function copyContextMeta( partIds: IdMap, ): void { if (!tableExists(db, "session_meta")) return; + // Native replay decisions were already filtered/remapped by copySessionStateForClone. const metaColumns = columns(db, "session_meta").filter( - (column) => column.name !== "session_id" && !RETIRED_META_COLUMNS.has(column.name), + (column) => + column.name !== "session_id" && + column.name !== "pi_native_tool_inputs" && + column.name !== "pi_native_reasoning_ids" && + !RETIRED_META_COLUMNS.has(column.name), ); const selectedColumns = ["session_id", ...metaColumns.map((column) => column.name)] .map(quoteIdentifier) diff --git a/packages/plugin/src/features/magic-context/storage-clone.ts b/packages/plugin/src/features/magic-context/storage-clone.ts index dac8e6927..cce9b3120 100644 --- a/packages/plugin/src/features/magic-context/storage-clone.ts +++ b/packages/plugin/src/features/magic-context/storage-clone.ts @@ -1,5 +1,6 @@ import { getHarness } from "../../shared/harness"; import type { Database } from "../../shared/sqlite"; +import { getNativeReasoningIds, getNativeToolInputs } from "./storage-native-replay"; export interface CloneCompartmentRow { sequence: number; @@ -220,6 +221,38 @@ function filterIdBlob(raw: string | null, filter: CloneSessionStateFilter): stri } } +function filterNativeToolInputs( + inputs: ReadonlyMap, + copiedToolCallIds: ReadonlySet, + filter: CloneSessionStateFilter, +): string { + const filtered = new Map(); + for (const [sourceId, serializedInput] of inputs) { + if (!copiedToolCallIds.has(sourceId)) continue; + const destinationId = mapMessageId(filter, sourceId); + if (destinationId === null) continue; + const existing = filtered.get(destinationId); + if (existing !== undefined && existing !== serializedInput) { + throw new Error(`native tool input clone collision for ${destinationId}`); + } + filtered.set(destinationId, serializedInput); + } + return JSON.stringify(Object.fromEntries(filtered)); +} + +function filterNativeReasoningIds( + ids: ReadonlySet, + filter: CloneSessionStateFilter, +): string { + const filtered = new Set(); + for (const sourceId of ids) { + if (!filter.includeMessageId(sourceId)) continue; + const destinationId = mapMessageId(filter, sourceId); + if (destinationId !== null) filtered.add(destinationId); + } + return JSON.stringify([...filtered]); +} + function clampWatermark(value: number | null, maxCopiedTag: number): number { if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) return 0; return Math.min(Math.floor(value), maxCopiedTag); @@ -323,6 +356,7 @@ export function copySessionStateForClone( ); const copiedTagNumbers: number[] = []; const copiedTagIds = new Map(); + const copiedToolCallIds = new Set(); for (const row of sourceTags) { if ( !filter.includeTag({ @@ -363,6 +397,7 @@ export function copySessionStateForClone( : sourceTagId; copiedTagIds.set(sourceTagId, destinationTagId); copiedTagNumbers.push(row.tag_number); + if (row.type === "tool") copiedToolCallIds.add(row.message_id); } if (copiedTagNumbers.length > 0) { @@ -491,6 +526,31 @@ export function copySessionStateForClone( migrateTodo ? (mapMessageId(filter, todoAnchor) ?? "") : "", migrateTodo ? (meta?.todo_synthetic_state_json ?? "") : "", ); + // The clone CLI also opens older databases without migrating their schema. + const metaColumns = new Set( + (db.prepare("PRAGMA table_info(session_meta)").all() as Array<{ name: string }>).map( + (column) => column.name, + ), + ); + if (metaColumns.has("pi_native_tool_inputs")) { + const inputs = filterNativeToolInputs( + getNativeToolInputs(db, sourceSessionId), + copiedToolCallIds, + filter, + ); + db.prepare( + "UPDATE session_meta SET pi_native_tool_inputs = ? WHERE session_id = ?", + ).run(inputs, destinationSessionId); + } + if (metaColumns.has("pi_native_reasoning_ids")) { + const ids = filterNativeReasoningIds( + getNativeReasoningIds(db, sourceSessionId), + filter, + ); + db.prepare( + "UPDATE session_meta SET pi_native_reasoning_ids = ? WHERE session_id = ?", + ).run(ids, destinationSessionId); + } const pendingOpsRow = db .prepare("SELECT COUNT(*) AS count FROM pending_ops WHERE session_id = ?") diff --git a/packages/plugin/src/features/magic-context/storage-db.ts b/packages/plugin/src/features/magic-context/storage-db.ts index 7b8502525..f58ecfb41 100644 --- a/packages/plugin/src/features/magic-context/storage-db.ts +++ b/packages/plugin/src/features/magic-context/storage-db.ts @@ -1816,6 +1816,8 @@ CREATE INDEX IF NOT EXISTS idx_dream_queue_pending ON dream_queue(started_at, en ensureColumn(db, "session_meta", "historian_last_failure_at", "INTEGER DEFAULT NULL"); ensureColumn(db, "session_meta", "system_prompt_hash", "TEXT DEFAULT ''"); ensureColumn(db, "session_meta", "cleared_reasoning_through_tag", "INTEGER DEFAULT 0"); + ensureColumn(db, "session_meta", "pi_native_tool_inputs", "TEXT DEFAULT NULL"); + ensureColumn(db, "session_meta", "pi_native_reasoning_ids", "TEXT NOT NULL DEFAULT '[]'"); ensureColumn(db, "session_meta", "tool_reclaim_watermark", "INTEGER DEFAULT 0"); ensureColumn(db, "session_meta", "stripped_placeholder_ids", "TEXT DEFAULT ''"); // Frozen replay watermark for the stale-ctx_reduce strip: message ids whose diff --git a/packages/plugin/src/features/magic-context/storage-native-replay.test.ts b/packages/plugin/src/features/magic-context/storage-native-replay.test.ts new file mode 100644 index 000000000..a09a01547 --- /dev/null +++ b/packages/plugin/src/features/magic-context/storage-native-replay.test.ts @@ -0,0 +1,185 @@ +/// + +import { afterEach, beforeEach, describe, expect, it } from "bun:test"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { Database } from "../../shared/sqlite"; +import { runMigrations } from "./migrations"; +import { initializeDatabase } from "./storage-db"; +import { + addNativeReasoningIds, + getNativeReasoningIds, + getNativeToolInputs, + saveNativeToolInputs, +} from "./storage-native-replay"; + +function createTestDb(): Database { + const db = new Database(":memory:"); + initializeDatabase(db); + runMigrations(db); + return db; +} + +describe("native replay storage", () => { + let db: Database; + + beforeEach(() => { + db = createTestDb(); + }); + + afterEach(() => { + db.close(); + }); + + it("treats absent and null native state as empty", () => { + expect(getNativeToolInputs(db, "missing")).toEqual(new Map()); + expect(getNativeReasoningIds(db, "missing")).toEqual(new Set()); + + const legacy = new Database(":memory:"); + try { + legacy.exec(` + CREATE TABLE session_meta ( + session_id TEXT PRIMARY KEY, + pi_native_tool_inputs TEXT, + pi_native_reasoning_ids TEXT + ); + `); + legacy + .prepare( + "INSERT INTO session_meta (session_id, pi_native_tool_inputs, pi_native_reasoning_ids) VALUES (?, NULL, NULL)", + ) + .run("legacy"); + + expect(getNativeToolInputs(legacy, "legacy")).toEqual(new Map()); + expect(getNativeReasoningIds(legacy, "legacy")).toEqual(new Set()); + } finally { + legacy.close(); + } + }); + + it("keeps legacy replay watermarks from seeding either native lane across an upgrade", () => { + db.prepare( + `INSERT INTO session_meta + (session_id, cleared_reasoning_through_tag, tool_reclaim_watermark, + stale_reduce_stripped_ids) + VALUES (?, ?, ?, ?)`, + ).run("legacy", 42, 24, JSON.stringify(["legacy-message"])); + db.exec("ALTER TABLE session_meta DROP COLUMN pi_native_tool_inputs"); + db.exec("ALTER TABLE session_meta DROP COLUMN pi_native_reasoning_ids"); + + initializeDatabase(db); + + expect(getNativeToolInputs(db, "legacy")).toEqual(new Map()); + expect(getNativeReasoningIds(db, "legacy")).toEqual(new Set()); + db.prepare("INSERT INTO session_meta (session_id) VALUES (?)").run("fresh"); + expect(getNativeToolInputs(db, "fresh")).toEqual(new Map()); + expect(getNativeReasoningIds(db, "fresh")).toEqual(new Set()); + }); + + it("persists only explicit native decisions and keeps the lanes independent", () => { + const initialInput = '{"path":"src/old.ts","marker":"[truncated]"}'; + const refreshedInput = '{"path":"src/old.ts","marker":"[full]"}'; + const secondInput = '{"path":"src/new.ts"}'; + + saveNativeToolInputs(db, "session", new Map([["call-1", initialInput]])); + + expect(getNativeToolInputs(db, "session")).toEqual(new Map([["call-1", initialInput]])); + expect(getNativeReasoningIds(db, "session")).toEqual(new Set()); + + addNativeReasoningIds(db, "session", ["assistant-1", "assistant-2"]); + saveNativeToolInputs( + db, + "session", + new Map([ + ["call-1", refreshedInput], + ["call-2", secondInput], + ]), + ); + + expect(getNativeToolInputs(db, "session")).toEqual( + new Map([ + ["call-1", refreshedInput], + ["call-2", secondInput], + ]), + ); + expect(getNativeReasoningIds(db, "session")).toEqual( + new Set(["assistant-1", "assistant-2"]), + ); + }); + + it("replays exact tool bytes and reasoning ids after reopening the database", () => { + const directory = mkdtempSync(join(tmpdir(), "magic-context-native-replay-")); + const path = join(directory, "context.db"); + const input = '{"path":"src/reopened.ts","range":{"start":1,"end":9}}'; + try { + const writer = new Database(path); + try { + initializeDatabase(writer); + runMigrations(writer); + saveNativeToolInputs(writer, "session", new Map([["call-1", input]])); + addNativeReasoningIds(writer, "session", ["assistant-1"]); + } finally { + writer.close(); + } + + const reader = new Database(path); + try { + expect(getNativeToolInputs(reader, "session")).toEqual( + new Map([["call-1", input]]), + ); + expect(getNativeReasoningIds(reader, "session")).toEqual(new Set(["assistant-1"])); + } finally { + reader.close(); + } + } finally { + rmSync(directory, { recursive: true, force: true }); + } + }); + + it("fails closed on malformed durable native state", () => { + db.prepare( + "INSERT INTO session_meta (session_id, pi_native_tool_inputs) VALUES (?, ?)", + ).run("bad-tool", '{"call-1":"not-json"}'); + db.prepare( + "INSERT INTO session_meta (session_id, pi_native_reasoning_ids) VALUES (?, ?)", + ).run("bad-reasoning", "not-json"); + + expect(() => getNativeToolInputs(db, "bad-tool")).toThrow( + "invalid persisted pi_native_tool_inputs state", + ); + expect(() => getNativeReasoningIds(db, "bad-reasoning")).toThrow( + "invalid persisted pi_native_reasoning_ids state", + ); + }); + + it("rolls back a failed write without changing either durable lane", () => { + const firstInput = '{"path":"src/first.ts"}'; + saveNativeToolInputs(db, "session", new Map([["call-1", firstInput]])); + addNativeReasoningIds(db, "session", ["assistant-1"]); + db.exec(` + CREATE TRIGGER fail_native_tool_input_update + BEFORE UPDATE OF pi_native_tool_inputs ON session_meta + WHEN NEW.session_id = 'session' + BEGIN + SELECT RAISE(ABORT, 'injected native tool write failure'); + END; + CREATE TRIGGER fail_native_reasoning_update + BEFORE UPDATE OF pi_native_reasoning_ids ON session_meta + WHEN NEW.session_id = 'session' + BEGIN + SELECT RAISE(ABORT, 'injected native reasoning write failure'); + END; + `); + + expect(() => + saveNativeToolInputs(db, "session", new Map([["call-2", '{"path":"src/second.ts"}']])), + ).toThrow("injected native tool write failure"); + expect(() => addNativeReasoningIds(db, "session", ["assistant-2"])).toThrow( + "injected native reasoning write failure", + ); + + expect(getNativeToolInputs(db, "session")).toEqual(new Map([["call-1", firstInput]])); + expect(getNativeReasoningIds(db, "session")).toEqual(new Set(["assistant-1"])); + }); +}); diff --git a/packages/plugin/src/features/magic-context/storage-native-replay.ts b/packages/plugin/src/features/magic-context/storage-native-replay.ts new file mode 100644 index 000000000..b4d1823ea --- /dev/null +++ b/packages/plugin/src/features/magic-context/storage-native-replay.ts @@ -0,0 +1,185 @@ +import { isRecord } from "../../shared/record-type-guard"; +import type { Database } from "../../shared/sqlite"; +import { ensureSessionMetaRow } from "./storage-meta-shared"; + +const NATIVE_TOOL_INPUTS_COLUMN = "pi_native_tool_inputs"; +const NATIVE_REASONING_IDS_COLUMN = "pi_native_reasoning_ids"; + +function invalidPersistedReplayState(column: string, sessionId: string): Error { + return new Error(`invalid persisted ${column} state for session ${sessionId}`); +} + +function assertNonEmptyId( + value: unknown, + column: string, + sessionId: string, +): asserts value is string { + if (typeof value !== "string" || value.length === 0) { + throw invalidPersistedReplayState(column, sessionId); + } +} + +function assertSerializedToolInput(value: unknown, sessionId: string): asserts value is string { + if (typeof value !== "string") { + throw invalidPersistedReplayState(NATIVE_TOOL_INPUTS_COLUMN, sessionId); + } + try { + if (!isRecord(JSON.parse(value))) { + throw new SyntaxError("native tool input must be an object"); + } + } catch { + throw invalidPersistedReplayState(NATIVE_TOOL_INPUTS_COLUMN, sessionId); + } +} + +function parseNativeToolInputs(raw: unknown, sessionId: string): Map { + if (raw === null || raw === undefined) return new Map(); + if (typeof raw !== "string") { + throw invalidPersistedReplayState(NATIVE_TOOL_INPUTS_COLUMN, sessionId); + } + + let parsed: unknown; + try { + parsed = JSON.parse(raw); + } catch { + throw invalidPersistedReplayState(NATIVE_TOOL_INPUTS_COLUMN, sessionId); + } + if (!isRecord(parsed)) { + throw invalidPersistedReplayState(NATIVE_TOOL_INPUTS_COLUMN, sessionId); + } + + const inputs = new Map(); + for (const [id, input] of Object.entries(parsed)) { + assertNonEmptyId(id, NATIVE_TOOL_INPUTS_COLUMN, sessionId); + assertSerializedToolInput(input, sessionId); + inputs.set(id, input); + } + return inputs; +} + +function parseNativeReasoningIds(raw: unknown, sessionId: string): Set { + if (raw === null || raw === undefined) return new Set(); + if (typeof raw !== "string") { + throw invalidPersistedReplayState(NATIVE_REASONING_IDS_COLUMN, sessionId); + } + + let parsed: unknown; + try { + parsed = JSON.parse(raw); + } catch { + throw invalidPersistedReplayState(NATIVE_REASONING_IDS_COLUMN, sessionId); + } + if (!Array.isArray(parsed)) { + throw invalidPersistedReplayState(NATIVE_REASONING_IDS_COLUMN, sessionId); + } + + const ids = new Set(); + for (const id of parsed) { + assertNonEmptyId(id, NATIVE_REASONING_IDS_COLUMN, sessionId); + ids.add(id); + } + return ids; +} + +/** + * Return frozen native tool inputs. Missing legacy state is empty; malformed + * stored state is rejected so replay never silently authorizes new bytes. + */ +export function getNativeToolInputs(db: Database, sessionId: string): Map { + const row = db + .prepare(`SELECT ${NATIVE_TOOL_INPUTS_COLUMN} FROM session_meta WHERE session_id = ?`) + .get(sessionId) as { pi_native_tool_inputs?: unknown } | undefined; + return parseNativeToolInputs(row?.pi_native_tool_inputs, sessionId); +} + +/** + * Atomically merge frozen native tool inputs. A supplied call id deliberately + * replaces its prior serialized input on an authorized cache-busting pass; + * values for every other call id remain intact. + */ +export function saveNativeToolInputs( + db: Database, + sessionId: string, + inputs: ReadonlyMap, +): void { + for (const [id, input] of inputs) { + assertNonEmptyId(id, NATIVE_TOOL_INPUTS_COLUMN, sessionId); + assertSerializedToolInput(input, sessionId); + } + + db.transaction(() => { + ensureSessionMetaRow(db, sessionId); + const row = db + .prepare(`SELECT ${NATIVE_TOOL_INPUTS_COLUMN} FROM session_meta WHERE session_id = ?`) + .get(sessionId) as { pi_native_tool_inputs?: unknown } | undefined; + const current = parseNativeToolInputs(row?.pi_native_tool_inputs, sessionId); + let changed = false; + for (const [id, input] of inputs) { + if (current.get(id) === input) continue; + current.set(id, input); + changed = true; + } + if (!changed) return; + + const result = db + .prepare( + `UPDATE session_meta SET ${NATIVE_TOOL_INPUTS_COLUMN} = ? WHERE session_id = ?`, + ) + .run(JSON.stringify(Object.fromEntries(current)), sessionId); + if (result.changes !== 1) { + throw new Error( + `failed to persist ${NATIVE_TOOL_INPUTS_COLUMN} for session ${sessionId}`, + ); + } + }).immediate(); +} + +/** + * Return assistant entries whose native reasoning was cleared. Missing legacy + * state is empty; malformed stored state fails closed. + */ +export function getNativeReasoningIds(db: Database, sessionId: string): Set { + const row = db + .prepare(`SELECT ${NATIVE_REASONING_IDS_COLUMN} FROM session_meta WHERE session_id = ?`) + .get(sessionId) as { pi_native_reasoning_ids?: unknown } | undefined; + return parseNativeReasoningIds(row?.pi_native_reasoning_ids, sessionId); +} + +/** Atomically union newly cleared native-reasoning entry ids into the replay set. */ +export function addNativeReasoningIds( + db: Database, + sessionId: string, + ids: Iterable, +): void { + const requested = new Set(); + for (const id of ids) { + assertNonEmptyId(id, NATIVE_REASONING_IDS_COLUMN, sessionId); + requested.add(id); + } + + db.transaction(() => { + ensureSessionMetaRow(db, sessionId); + const row = db + .prepare(`SELECT ${NATIVE_REASONING_IDS_COLUMN} FROM session_meta WHERE session_id = ?`) + .get(sessionId) as { pi_native_reasoning_ids?: unknown } | undefined; + const current = parseNativeReasoningIds(row?.pi_native_reasoning_ids, sessionId); + let changed = false; + for (const id of requested) { + if (current.has(id)) continue; + current.add(id); + changed = true; + } + if (!changed) return; + + const result = db + .prepare( + `UPDATE session_meta SET ${NATIVE_REASONING_IDS_COLUMN} = ? WHERE session_id = ?`, + ) + .run(JSON.stringify([...current]), sessionId); + if (result.changes !== 1) { + throw new Error( + `failed to persist ${NATIVE_REASONING_IDS_COLUMN} for session ${sessionId}`, + ); + } + }).immediate(); +} From 88b1a59a74399545d9148d4e1138aa7943e0bef1 Mon Sep 17 00:00:00 2001 From: SilentBless <126187696+SilentBless@users.noreply.github.com> Date: Fri, 11 Sep 2026 05:47:21 +0300 Subject: [PATCH 5/7] fix(pi): retain native history on corrupt replay state --- packages/pi-plugin/README.md | 2 + packages/pi-plugin/src/context-handler.ts | 66 +++++-- .../src/native-replay-state-pi.test.ts | 185 ++++++++++++------ .../pi-plugin/src/native-replay-state-pi.ts | 6 +- 4 files changed, 178 insertions(+), 81 deletions(-) diff --git a/packages/pi-plugin/README.md b/packages/pi-plugin/README.md index 89f7a338f..5028e756a 100644 --- a/packages/pi-plugin/README.md +++ b/packages/pi-plugin/README.md @@ -80,6 +80,8 @@ OMP can replay Responses/Codex history from `providerPayload` instead of ordinar Native tool-input values and native reasoning removals have separate persisted replay state. Existing dropped tags and local reasoning watermarks do not activate native changes during upgrade: first application waits for an already-authorized cache-busting pass and is persisted before publication. Deferred passes replay only those saved native decisions. +If either native replay field cannot be read or validated, that pass skips native replay and activation, preserving the incoming native payloads. Ordinary Pi cleanup continues; the stored native decisions are not reset. + Function calls carry the canonical dropped-marker JSON. Custom calls use OMP's existing empty-string fallback, which does not carry that marker or imply the same copied-input rejection behavior. Display summaries (`summary`) are removed with eligible old encrypted reasoning; they are the source of ordinary Pi `thinking`, not a native preservation requirement. Retaining a native payload does not prevent per-part cleanup of stale non-redacted Pi thinking and its signature. diff --git a/packages/pi-plugin/src/context-handler.ts b/packages/pi-plugin/src/context-handler.ts index 414566201..918f12a94 100644 --- a/packages/pi-plugin/src/context-handler.ts +++ b/packages/pi-plugin/src/context-handler.ts @@ -107,6 +107,10 @@ import { pruneAutoSearchHintDecisions, pruneNoteNudgeAnchors, } from "@magic-context/core/features/magic-context/storage-meta-persisted"; +import { + getNativeReasoningIds, + getNativeToolInputs, +} from "@magic-context/core/features/magic-context/storage-native-replay"; import { getSourceContents } from "@magic-context/core/features/magic-context/storage-source"; import { createTagger, @@ -5831,30 +5835,50 @@ async function runPipeline(args: RunPipelineArgs): Promise { // Legacy drop/watermark state does not authorize first native activation. // Use committed canonical inputs, then publish only persisted native decisions. - const nativeInputsApplied = applyNativeToolInputReplayPi({ - db: args.db, - sessionId: args.sessionId, - messages: args.messages, - changes: transcript.getToolInputChanges(), - canApply: isCacheBustingPass, - }); - const nativeReasoningApplied = args.reasoningClearing - ? applyNativeReasoningReplayPi({ + let nativeInputs: ReadonlyMap | undefined; + let nativeReasoningIds: ReadonlySet | undefined; + try { + // Validate both lanes before either can publish replay or activation. + nativeInputs = getNativeToolInputs(args.db, args.sessionId); + nativeReasoningIds = getNativeReasoningIds(args.db, args.sessionId); + } catch (error) { + sessionLog( + args.sessionId, + `native replay state unavailable; retaining native history (continuing): ${error instanceof Error ? error.message : String(error)}`, + ); + } + if (nativeInputs !== undefined && nativeReasoningIds !== undefined) { + const nativeInputsApplied = applyNativeToolInputReplayPi( + { db: args.db, sessionId: args.sessionId, messages: args.messages, - messageIdToMaxTag, - stableId: stableIdResolver, - localWatermark: args.sessionMeta.clearedReasoningThroughTag ?? 0, - clearReasoningAge: args.reasoningClearing.clearReasoningAge, - omissionAllowed: args.reasoningClearing.nativeReasoningMayClear, - canApply: isCacheBustingPass && !reasoningPersistenceFailed, - detectAged: shouldRunHeuristics && routineCleanupApplied, - }) - : 0; - if (nativeInputsApplied > 0 || nativeReasoningApplied > 0) { - heuristicOrReasoningDidMutate = true; - executedWorkThisPass = true; + changes: transcript.getToolInputChanges(), + canApply: isCacheBustingPass, + }, + nativeInputs, + ); + const nativeReasoningApplied = args.reasoningClearing + ? applyNativeReasoningReplayPi( + { + db: args.db, + sessionId: args.sessionId, + messages: args.messages, + messageIdToMaxTag, + stableId: stableIdResolver, + localWatermark: args.sessionMeta.clearedReasoningThroughTag ?? 0, + clearReasoningAge: args.reasoningClearing.clearReasoningAge, + omissionAllowed: args.reasoningClearing.nativeReasoningMayClear, + canApply: isCacheBustingPass && !reasoningPersistenceFailed, + detectAged: shouldRunHeuristics && routineCleanupApplied, + }, + nativeReasoningIds, + ) + : 0; + if (nativeInputsApplied > 0 || nativeReasoningApplied > 0) { + heuristicOrReasoningDidMutate = true; + executedWorkThisPass = true; + } } if (toolReclaimApplicationOpportunity) { advanceToolReclaimWatermarkToCurrentMax(args.db, args.sessionId); diff --git a/packages/pi-plugin/src/native-replay-state-pi.test.ts b/packages/pi-plugin/src/native-replay-state-pi.test.ts index dcd7184be..174867a70 100644 --- a/packages/pi-plugin/src/native-replay-state-pi.test.ts +++ b/packages/pi-plugin/src/native-replay-state-pi.test.ts @@ -219,6 +219,64 @@ describe("native upgrade application", () => { } }); + for (const column of [ + "pi_native_tool_inputs", + "pi_native_reasoning_ids", + ] as const) { + it(`preserves both native lanes and continues local cleanup when ${column} is malformed`, async () => { + const sessionId = `ses-native-malformed-${column}`; + const f = fixture(sessionId); + try { + await f.seedLegacy(); + const stored = { + pi_native_tool_inputs: + column === "pi_native_tool_inputs" + ? "{" + : JSON.stringify({ + [callId]: JSON.stringify({ + dropped: "persisted native marker", + }), + }), + pi_native_reasoning_ids: + column === "pi_native_reasoning_ids" + ? "{}" + : JSON.stringify(["entry-old"]), + }; + f.db + .prepare( + "UPDATE session_meta SET pi_native_tool_inputs = ?, pi_native_reasoning_ids = ? WHERE session_id = ?", + ) + .run( + stored.pi_native_tool_inputs, + stored.pi_native_reasoning_ids, + sessionId, + ); + const before = nativeBytes([oldAssistant()]); + for (const percent of [0, 90, 0]) { + f.restart(); + const messages = await f.pass(percent); + expect(nativeBytes(messages)).toBe(before); + expect(messages[1]).toMatchObject({ + content: [ + { type: "thinking", thinking: "", thinkingSignature: undefined }, + {}, + { arguments: { dropped: expect.any(String) } }, + ], + }); + } + expect( + f.db + .prepare( + "SELECT pi_native_tool_inputs, pi_native_reasoning_ids FROM session_meta WHERE session_id = ?", + ) + .get(sessionId), + ).toEqual(stored); + } finally { + f.close(); + } + }); + } + for (const column of [ "pi_native_tool_inputs", "pi_native_reasoning_ids", @@ -258,13 +316,16 @@ describe("native upgrade application", () => { }; message.content[2].arguments = { dropped: marker }; const messages: unknown[] = [message]; - applyNativeToolInputReplayPi({ - db, - sessionId, - messages, - canApply, - changes: changed ? new Map([[0, new Set([callId])]]) : new Map(), - }); + applyNativeToolInputReplayPi( + { + db, + sessionId, + messages, + canApply, + changes: changed ? new Map([[0, new Set([callId])]]) : new Map(), + }, + getNativeToolInputs(db, sessionId), + ); return nativeBytes(messages); }; const b = run("first marker", true); @@ -291,18 +352,21 @@ describe("native upgrade application", () => { try { const run = (id: string, canApply: boolean) => { const messages: unknown[] = [oldAssistant()]; - applyNativeReasoningReplayPi({ - db, - sessionId, - messages, - stableId: () => id, - messageIdToMaxTag: new Map([[id, 1]]), - localWatermark: 10, - clearReasoningAge: 100, - omissionAllowed: true, - detectAged: false, - canApply, - }); + applyNativeReasoningReplayPi( + { + db, + sessionId, + messages, + stableId: () => id, + messageIdToMaxTag: new Map([[id, 1]]), + localWatermark: 10, + clearReasoningAge: 100, + omissionAllowed: true, + detectAged: false, + canApply, + }, + getNativeReasoningIds(db, sessionId), + ); return nativeBytes(messages); }; const unresolved = run("pi-msg-0-2-assistant", true); @@ -327,21 +391,24 @@ describe("native upgrade application", () => { const before = nativeBytes([original]); const run = (canApply: boolean) => { const messages = [original]; - applyNativeReasoningReplayPi({ - db, - sessionId, - messages, - messageIdToMaxTag: new Map([ - ["entry-old", 1], - ["entry-new", 100], - ]), - stableId: () => "entry-old", - localWatermark: 0, - clearReasoningAge: 10, - omissionAllowed: true, - detectAged: canApply, - canApply, - }); + applyNativeReasoningReplayPi( + { + db, + sessionId, + messages, + messageIdToMaxTag: new Map([ + ["entry-old", 1], + ["entry-new", 100], + ]), + stableId: () => "entry-old", + localWatermark: 0, + clearReasoningAge: 10, + omissionAllowed: true, + detectAged: canApply, + canApply, + }, + getNativeReasoningIds(db, sessionId), + ); return nativeBytes(messages); }; const after = run(true); @@ -366,13 +433,16 @@ describe("native upgrade application", () => { content: Array<{ arguments?: Record }>; }; toolMessage.content[2].arguments = { dropped: "native tool marker" }; - applyNativeToolInputReplayPi({ - db, - sessionId, - messages: [toolMessage], - changes: new Map([[0, new Set([callId])]]), - canApply: true, - }); + applyNativeToolInputReplayPi( + { + db, + sessionId, + messages: [toolMessage], + changes: new Map([[0, new Set([callId])]]), + canApply: true, + }, + getNativeToolInputs(db, sessionId), + ); expect(getNativeToolInputs(db, sessionId).has(callId)).toBe(true); const run = ( omissionAllowed: boolean, @@ -380,21 +450,24 @@ describe("native upgrade application", () => { localWatermark: number, ) => { const messages: unknown[] = [oldAssistant()]; - applyNativeReasoningReplayPi({ - db, - sessionId, - messages, - messageIdToMaxTag: new Map([ - ["old", 1], - ["new", 200], - ]), - stableId: () => "old", - localWatermark, - clearReasoningAge: 10, - omissionAllowed, - canApply, - detectAged: false, - }); + applyNativeReasoningReplayPi( + { + db, + sessionId, + messages, + messageIdToMaxTag: new Map([ + ["old", 1], + ["new", 200], + ]), + stableId: () => "old", + localWatermark, + clearReasoningAge: 10, + omissionAllowed, + canApply, + detectAged: false, + }, + getNativeReasoningIds(db, sessionId), + ); return nativeBytes(messages); }; expect(run(false, true, 10)).toContain(ciphertext); diff --git a/packages/pi-plugin/src/native-replay-state-pi.ts b/packages/pi-plugin/src/native-replay-state-pi.ts index c2ca900fe..dfdeddb4b 100644 --- a/packages/pi-plugin/src/native-replay-state-pi.ts +++ b/packages/pi-plugin/src/native-replay-state-pi.ts @@ -1,8 +1,6 @@ import type { ContextDatabase } from "@magic-context/core/features/magic-context/storage"; import { addNativeReasoningIds, - getNativeReasoningIds, - getNativeToolInputs, saveNativeToolInputs, } from "@magic-context/core/features/magic-context/storage-native-replay"; import { sessionLog } from "@magic-context/core/shared/logger"; @@ -25,8 +23,8 @@ export function applyNativeToolInputReplayPi( changes: ReadonlyMap>; canApply: boolean; }, + saved: ReadonlyMap, ): number { - const saved = getNativeToolInputs(args.db, args.sessionId); if (saved.size === 0 && (!args.canApply || args.changes.size === 0)) return 0; const nextInputs = new Map(); const pending = new Map>(); @@ -110,9 +108,9 @@ export function applyNativeReasoningReplayPi( canApply: boolean; detectAged: boolean; }, + saved: ReadonlySet, ): number { if (!args.omissionAllowed) return 0; - const saved = getNativeReasoningIds(args.db, args.sessionId); let maxTag = 0; for (const tag of args.messageIdToMaxTag.values()) maxTag = Math.max(maxTag, tag); From 29d0b26a2d458d435eef3d726aa90ef73ef3037b Mon Sep 17 00:00:00 2001 From: SilentBless <126187696+SilentBless@users.noreply.github.com> Date: Fri, 11 Sep 2026 05:47:21 +0300 Subject: [PATCH 6/7] refactor: clarify native replay ownership and reuse storage writes --- packages/pi-plugin/README.md | 12 ---- .../src/native-replay-state-pi.test.ts | 43 +++++++++++++- .../magic-context/storage-native-replay.ts | 56 +++++++++---------- 3 files changed, 69 insertions(+), 42 deletions(-) diff --git a/packages/pi-plugin/README.md b/packages/pi-plugin/README.md index 5028e756a..ed8e9e549 100644 --- a/packages/pi-plugin/README.md +++ b/packages/pi-plugin/README.md @@ -76,18 +76,6 @@ npx @cortexkit/magic-context@latest doctor --harness omp OMP's legacy Pi loader maps `@earendil-works/*` imports to its bundled `@oh-my-pi/*` runtime. Magic Context identifies OMP from `@oh-my-pi/pi-utils`'s in-process `APP_NAME` through the running host's module graph, not from an executable basename, so compiled and symlinked launches behave the same. It also re-invokes the current host executable for historian, dreamer, and sidekick children, so OMP children remain OMP processes. -OMP can replay Responses/Codex history from `providerPayload` instead of ordinary message content. Tool-input reductions update the matched native call without discarding unrelated native results. Old encrypted reasoning is cleared only for Codex models whose resolved compatibility settings explicitly allow omission, and only from incremental history. Other models, full snapshots, plaintext or malformed reasoning, redacted thinking, and computer-linked reasoning retain their native reasoning. Native text and user/developer history carriers are not rewritten. - -Native tool-input values and native reasoning removals have separate persisted replay state. Existing dropped tags and local reasoning watermarks do not activate native changes during upgrade: first application waits for an already-authorized cache-busting pass and is persisted before publication. Deferred passes replay only those saved native decisions. - -If either native replay field cannot be read or validated, that pass skips native replay and activation, preserving the incoming native payloads. Ordinary Pi cleanup continues; the stored native decisions are not reset. - -Function calls carry the canonical dropped-marker JSON. Custom calls use OMP's existing empty-string fallback, which does not carry that marker or imply the same copied-input rejection behavior. - -Display summaries (`summary`) are removed with eligible old encrypted reasoning; they are the source of ordinary Pi `thinking`, not a native preservation requirement. Retaining a native payload does not prevent per-part cleanup of stale non-redacted Pi thinking and its signature. - -OMP exposes no dedicated native-item mutation API, so this adapter targets its current Responses history representation. It changes the request transcript; it does not perform a stored session-JSONL rewrite. - --- ## Configuration diff --git a/packages/pi-plugin/src/native-replay-state-pi.test.ts b/packages/pi-plugin/src/native-replay-state-pi.test.ts index 174867a70..4f4ba0c0b 100644 --- a/packages/pi-plugin/src/native-replay-state-pi.test.ts +++ b/packages/pi-plugin/src/native-replay-state-pi.test.ts @@ -115,9 +115,9 @@ function fixture(sessionId: string) { heuristics: { clearReasoningAge: 100 }, scheduler: { executeThresholdPercentage: 80 }, }; - const restart = () => { + const restart = (compactionOff = false) => { clearContextHandlerSession(sessionId); - registerPiContextHandler(fake.pi as never, options); + registerPiContextHandler(fake.pi as never, { ...options, compactionOff }); }; restart(); const pass = async (percent: number) => { @@ -219,6 +219,45 @@ describe("native upgrade application", () => { } }); + it("suspends saved native decisions in compaction-off mode and resumes them when enabled", async () => { + const sessionId = "ses-native-compaction-off"; + const f = fixture(sessionId); + try { + await f.seedLegacy(); + const reduced = nativeBytes(await f.pass(90)); + expect(reduced).not.toContain(staleInput); + expect(reduced).not.toContain(ciphertext); + const savedInputs = getNativeToolInputs(f.db, sessionId); + const savedReasoning = getNativeReasoningIds(f.db, sessionId); + const original = nativeBytes([oldAssistant()]); + + for (const percent of [0, 90, 0]) { + f.restart(true); + const messages = await f.pass(percent); + expect(nativeBytes(messages)).toBe(original); + expect(messages[1]).toMatchObject({ + content: [ + { + type: "thinking", + thinking: "old summary", + thinkingSignature: "old signature", + }, + { type: "text", text: "visible answer" }, + { arguments: { path: staleInput } }, + ], + }); + expect(getNativeToolInputs(f.db, sessionId)).toEqual(savedInputs); + expect(getNativeReasoningIds(f.db, sessionId)).toEqual(savedReasoning); + } + + f.restart(); + expect(nativeBytes(await f.pass(0))).toBe(reduced); + expect(nativeBytes(await f.pass(0))).toBe(reduced); + } finally { + f.close(); + } + }); + for (const column of [ "pi_native_tool_inputs", "pi_native_reasoning_ids", diff --git a/packages/plugin/src/features/magic-context/storage-native-replay.ts b/packages/plugin/src/features/magic-context/storage-native-replay.ts index b4d1823ea..37a54029d 100644 --- a/packages/plugin/src/features/magic-context/storage-native-replay.ts +++ b/packages/plugin/src/features/magic-context/storage-native-replay.ts @@ -81,6 +81,20 @@ function parseNativeReasoningIds(raw: unknown, sessionId: string): Set { return ids; } +function writeNativeReplayState( + db: Database, + sessionId: string, + column: typeof NATIVE_TOOL_INPUTS_COLUMN | typeof NATIVE_REASONING_IDS_COLUMN, + serialized: string, +): void { + const result = db + .prepare(`UPDATE session_meta SET ${column} = ? WHERE session_id = ?`) + .run(serialized, sessionId); + if (result.changes !== 1) { + throw new Error(`failed to persist ${column} for session ${sessionId}`); + } +} + /** * Return frozen native tool inputs. Missing legacy state is empty; malformed * stored state is rejected so replay never silently authorizes new bytes. @@ -109,10 +123,7 @@ export function saveNativeToolInputs( db.transaction(() => { ensureSessionMetaRow(db, sessionId); - const row = db - .prepare(`SELECT ${NATIVE_TOOL_INPUTS_COLUMN} FROM session_meta WHERE session_id = ?`) - .get(sessionId) as { pi_native_tool_inputs?: unknown } | undefined; - const current = parseNativeToolInputs(row?.pi_native_tool_inputs, sessionId); + const current = getNativeToolInputs(db, sessionId); let changed = false; for (const [id, input] of inputs) { if (current.get(id) === input) continue; @@ -121,16 +132,12 @@ export function saveNativeToolInputs( } if (!changed) return; - const result = db - .prepare( - `UPDATE session_meta SET ${NATIVE_TOOL_INPUTS_COLUMN} = ? WHERE session_id = ?`, - ) - .run(JSON.stringify(Object.fromEntries(current)), sessionId); - if (result.changes !== 1) { - throw new Error( - `failed to persist ${NATIVE_TOOL_INPUTS_COLUMN} for session ${sessionId}`, - ); - } + writeNativeReplayState( + db, + sessionId, + NATIVE_TOOL_INPUTS_COLUMN, + JSON.stringify(Object.fromEntries(current)), + ); }).immediate(); } @@ -159,10 +166,7 @@ export function addNativeReasoningIds( db.transaction(() => { ensureSessionMetaRow(db, sessionId); - const row = db - .prepare(`SELECT ${NATIVE_REASONING_IDS_COLUMN} FROM session_meta WHERE session_id = ?`) - .get(sessionId) as { pi_native_reasoning_ids?: unknown } | undefined; - const current = parseNativeReasoningIds(row?.pi_native_reasoning_ids, sessionId); + const current = getNativeReasoningIds(db, sessionId); let changed = false; for (const id of requested) { if (current.has(id)) continue; @@ -171,15 +175,11 @@ export function addNativeReasoningIds( } if (!changed) return; - const result = db - .prepare( - `UPDATE session_meta SET ${NATIVE_REASONING_IDS_COLUMN} = ? WHERE session_id = ?`, - ) - .run(JSON.stringify([...current]), sessionId); - if (result.changes !== 1) { - throw new Error( - `failed to persist ${NATIVE_REASONING_IDS_COLUMN} for session ${sessionId}`, - ); - } + writeNativeReplayState( + db, + sessionId, + NATIVE_REASONING_IDS_COLUMN, + JSON.stringify([...current]), + ); }).immediate(); } From c27fc840ebc21c7e5beea2fc2643f6138f262aff Mon Sep 17 00:00:00 2001 From: SilentBless <126187696+SilentBless@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:30:11 +0300 Subject: [PATCH 7/7] fix: remap trailing replay decisions when cloning --- .../pi-plugin/src/clone-inheritance.test.ts | 20 +++++++++++++++++++ packages/plugin/scripts/clone-session.test.ts | 3 +-- .../features/magic-context/storage-clone.ts | 12 +++++++++++ 3 files changed, 33 insertions(+), 2 deletions(-) diff --git a/packages/pi-plugin/src/clone-inheritance.test.ts b/packages/pi-plugin/src/clone-inheritance.test.ts index 5aec6f182..23d8a8f5d 100644 --- a/packages/pi-plugin/src/clone-inheritance.test.ts +++ b/packages/pi-plugin/src/clone-inheritance.test.ts @@ -819,11 +819,31 @@ describe("Pi clone state inheritance", () => { new Set(["clone-assistant-retained"]), ); expect(getTrailingBlankDecisions(database, "clone")).toEqual( + new Map([["clone-assistant-retained", "strip"]]), + ); + }); + + it("filters and remaps legacy flat trailing decisions without adding native state", () => { + const database = db(); + addTrailingBlankDecisions(database, "source", [ + ["assistant-retained", "strip"], + ["assistant-outside", "keep:2"], + ]); + copySessionStateForClone(database, "source", "clone", { + ...__test.createCloneFilter([assistant("assistant-retained")]), + mapMessageId: (id) => `clone-${id}`, + }); + expect(getTrailingBlankDecisions(database, "clone")).toEqual( + new Map([["clone-assistant-retained", "strip"]]), + ); + expect(getTrailingBlankDecisions(database, "source")).toEqual( new Map([ ["assistant-retained", "strip"], ["assistant-outside", "keep:2"], ]), ); + expect(getNativeToolInputs(database, "clone")).toEqual(new Map()); + expect(getNativeReasoningIds(database, "clone")).toEqual(new Set()); }); it("fails closed and rolls back a clone with malformed native replay", () => { diff --git a/packages/plugin/scripts/clone-session.test.ts b/packages/plugin/scripts/clone-session.test.ts index 23107736a..045e87c9d 100644 --- a/packages/plugin/scripts/clone-session.test.ts +++ b/packages/plugin/scripts/clone-session.test.ts @@ -564,8 +564,7 @@ describe("clone-session", () => { const replayDocument = parseReplayDocument(state.trailing_blank_decisions); expect(replayDocument.version).toBe(2); expect(Object.entries(replayDocument.trailingBlank)).toEqual([ - ["msg_source_2", "strip"], - ["outside-assistant", "keep:2"], + [tool.tool_owner_message_id, "strip"], ]); expect(replayDocument.piNative).toEqual({ toolInputs: { [tool.message_id]: frozenInput }, diff --git a/packages/plugin/src/features/magic-context/storage-clone.ts b/packages/plugin/src/features/magic-context/storage-clone.ts index c815b280d..26907d009 100644 --- a/packages/plugin/src/features/magic-context/storage-clone.ts +++ b/packages/plugin/src/features/magic-context/storage-clone.ts @@ -289,6 +289,18 @@ function cloneReplayDocument( filter: CloneSessionStateFilter, ): ReplayDocument { const document = readReplayDocument(db, sourceSessionId); + const trailingBlank = new Map(); + for (const [sourceId, decision] of Object.entries(document.trailingBlank)) { + if (!filter.includeMessageId(sourceId)) continue; + const destinationId = mapMessageId(filter, sourceId); + if (destinationId === null) continue; + const existing = trailingBlank.get(destinationId); + if (existing !== undefined && existing !== decision) { + throw new Error(`trailing blank clone collision for ${destinationId}`); + } + trailingBlank.set(destinationId, decision); + } + document.trailingBlank = Object.fromEntries(trailingBlank); if (document.version === 1 || document.piNative === undefined) return document; const nativeReplay = getNativeReplayState(db, sourceSessionId);