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.test.ts b/packages/pi-plugin/src/context-handler.test.ts index 603c52c8a..a20dfc3aa 100644 --- a/packages/pi-plugin/src/context-handler.test.ts +++ b/packages/pi-plugin/src/context-handler.test.ts @@ -3619,11 +3619,83 @@ 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"; + let watermarkWriteAttempted = false; const restorePersistence = contextHandlerInternals.setReasoningWatermarkPersistenceForTests(() => { + watermarkWriteAttempted = true; throw new Error("faulted reasoning watermark write"); }); try { @@ -3642,6 +3714,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 +3746,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 +3767,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..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, @@ -223,6 +227,11 @@ import { prepareCachedM0M1PiReplay, 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, @@ -3086,6 +3095,7 @@ export function registerPiContextHandler( clearReasoningAge: options.heuristics?.clearReasoningAge ?? DEFAULT_CLEAR_REASONING_AGE, + nativeReasoningMayClear: canClearNativeReasoning(ctx.model), }, canUseEmptySentinels, temporalAwareness: options.injection?.temporalAwareness === true, @@ -4529,6 +4539,7 @@ interface RunPipelineArgs { */ reasoningClearing?: { clearReasoningAge: number; + nativeReasoningMayClear: boolean; }; /** True only when the active provider filters empty sentinel content safely. */ canUseEmptySentinels: boolean; @@ -4631,7 +4642,10 @@ function captureReasoningMutationRollback( }> = []; 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; + }; if (message.role !== "assistant" || !Array.isArray(message.content)) continue; for (const rawPart of message.content) { @@ -5642,6 +5656,7 @@ 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); try { @@ -5687,6 +5702,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. @@ -5816,6 +5832,54 @@ 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. + 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, + 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-pi.test.ts b/packages/pi-plugin/src/native-replay-pi.test.ts new file mode 100644 index 000000000..905869b60 --- /dev/null +++ b/packages/pi-plugin/src/native-replay-pi.test.ts @@ -0,0 +1,364 @@ +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("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", + 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("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" }, + ]); + 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..44b7a060c --- /dev/null +++ b/packages/pi-plugin/src/native-replay-pi.ts @@ -0,0 +1,191 @@ +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"; + 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(); + 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/native-replay-state-pi.test.ts b/packages/pi-plugin/src/native-replay-state-pi.test.ts new file mode 100644 index 000000000..4f4ba0c0b --- /dev/null +++ b/packages/pi-plugin/src/native-replay-state-pi.test.ts @@ -0,0 +1,523 @@ +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 = (compactionOff = false) => { + clearContextHandlerSession(sessionId); + registerPiContextHandler(fake.pi as never, { ...options, compactionOff }); + }; + 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(); + } + }); + + 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", + ] 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", + ] 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(), + }, + getNativeToolInputs(db, sessionId), + ); + 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, + }, + getNativeReasoningIds(db, sessionId), + ); + 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, + }, + getNativeReasoningIds(db, sessionId), + ); + 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, + }, + getNativeToolInputs(db, sessionId), + ); + 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, + }, + getNativeReasoningIds(db, sessionId), + ); + 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..dfdeddb4b --- /dev/null +++ b/packages/pi-plugin/src/native-replay-state-pi.ts @@ -0,0 +1,155 @@ +import type { ContextDatabase } from "@magic-context/core/features/magic-context/storage"; +import { + addNativeReasoningIds, + 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; + }, + saved: ReadonlyMap, +): number { + 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; + }, + saved: ReadonlySet, +): number { + if (!args.omissionAllowed) return 0; + 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 f0f443da5..067ac22ce 100644 --- a/packages/pi-plugin/src/reasoning-replay-pi.test.ts +++ b/packages/pi-plugin/src/reasoning-replay-pi.test.ts @@ -128,6 +128,118 @@ describe("clearOldReasoningPi", () => { }); }); + 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", + }); + + 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", () => { + const db = makeDb(); + const sessionId = "ses-native-snapshot"; + try { + const native = { + type: "openaiResponsesHistory", + dt: false, + items: [{ type: "reasoning", encrypted_content: "snapshot-reasoning" }], + }; + const original = { + role: "assistant", + content: [ + { + type: "thinking", + thinking: "required history", + thinkingSignature: "sig", + }, + { type: "text", text: "reply" }, + ], + providerPayload: native, + }; + const nativeSnapshot = structuredClone(native); + const messageIdToMaxTag = new Map([ + ["a", 1], + ["recent", 10], + ]); + const options = { + messageIdToMaxTag, + piMessageStableId: () => "a", + }; + + 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("does nothing when ageCutoff is 0 or below", () => { const messages = [ { @@ -149,15 +261,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", @@ -165,25 +285,59 @@ 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, + }; + + 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 f30af2cfd..a042d17c5 100644 --- a/packages/pi-plugin/src/reasoning-replay-pi.ts +++ b/packages/pi-plugin/src/reasoning-replay-pi.ts @@ -109,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). */ @@ -146,6 +146,7 @@ export function clearOldReasoningPi(args: { const msgTag = messageIdToMaxTag.get(id) ?? 0; if (msgTag === 0 || msgTag > ageCutoff) continue; + let clearedThisMessage = false; for (const part of msg.content) { if ( part && @@ -170,11 +171,13 @@ export function clearOldReasoningPi(args: { tp.thinking = CLEARED; tp.thinkingSignature = undefined; cleared++; + clearedThisMessage = true; } } } - - if (cleared > 0 && msgTag > newWatermark) newWatermark = msgTag; + if (clearedThisMessage && msgTag > newWatermark) { + newWatermark = msgTag; + } } return { cleared, newWatermark }; @@ -241,10 +244,9 @@ 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; 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..621f81762 100644 --- a/packages/pi-plugin/src/transcript-pi.test.ts +++ b/packages/pi-plugin/src/transcript-pi.test.ts @@ -75,6 +75,138 @@ describe("createPiTranscript", () => { expect(transcript.getOutputMessages()).toBe(messages); }); + 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: "toolCall", + id: "call-sentinel|fc-sentinel", + name: "write", + arguments: { 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" }, + }, + ], + }), + 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: { 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", () => { const vectors = [ { diff --git a/packages/pi-plugin/src/transcript-pi.ts b/packages/pi-plugin/src/transcript-pi.ts index 5b62df5a6..c2f43a666 100644 --- a/packages/pi-plugin/src/transcript-pi.ts +++ b/packages/pi-plugin/src/transcript-pi.ts @@ -132,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 @@ -163,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 @@ -174,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 @@ -184,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, ); @@ -227,6 +242,9 @@ export function createPiTranscript( getWorkingMessages(): PiAgentMessage[] { return working; }, + getToolInputChanges(): ReadonlyMap> { + return toolInputChanges; + }, }; } @@ -240,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[] = []; @@ -342,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; @@ -404,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 { @@ -449,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; @@ -471,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]; @@ -499,7 +517,7 @@ function createOpaqueTranscriptMessage( function createPiUserStringPart( working: PiAgentMessage[], messageIndex: number, - markDirty: (messageIndex: number) => void, + markDirty: MarkDirty, ): TranscriptPart { return { kind: "text", @@ -543,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; @@ -604,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]; @@ -674,7 +692,7 @@ function createPiAssistantPart( ...(working[messageIndex] as PiAssistantMessage), content: newContent, }; - markDirty(messageIndex); + markDirty(messageIndex, p.id); return true; } return false; @@ -733,7 +751,7 @@ function createPiAssistantPart( ...(working[messageIndex] as PiAssistantMessage), content: newContent, }; - markDirty(messageIndex); + markDirty(messageIndex, p.id); return true; }, // Replace this assistant part's content with a sentinel placeholder. @@ -777,7 +795,10 @@ function createPiAssistantPart( ...(working[messageIndex] as PiAssistantMessage), content: newContent, }; - markDirty(messageIndex); + markDirty( + messageIndex, + existing?.type === "toolCall" ? existing.id : undefined, + ); return true; }, }; @@ -787,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..37a54029d --- /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; +} + +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. + */ +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 current = getNativeToolInputs(db, 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; + + writeNativeReplayState( + db, + sessionId, + NATIVE_TOOL_INPUTS_COLUMN, + JSON.stringify(Object.fromEntries(current)), + ); + }).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 current = getNativeReasoningIds(db, sessionId); + let changed = false; + for (const id of requested) { + if (current.has(id)) continue; + current.add(id); + changed = true; + } + if (!changed) return; + + writeNativeReplayState( + db, + sessionId, + NATIVE_REASONING_IDS_COLUMN, + JSON.stringify([...current]), + ); + }).immediate(); +}