From f7344f662e79d739c2c97335f8b95b18fd252172 Mon Sep 17 00:00:00 2001 From: Omri Hazan Date: Wed, 22 Jul 2026 12:29:54 +0300 Subject: [PATCH] fix: reconstruct compression state after fork and compaction When a session is forked or compacted, DCP loses all compression state because the new session has no persisted state file (fork) or state is wiped (compaction). This causes previously compressed messages to leak back into context, blowing up the context window. This fix detects sessions that have compress tool results in their conversation history but no persisted compression state, and reconstructs the compression blocks by replaying the compress tool results in chronological order. Fixes #457, #521 --- lib/compress/reconstruct.ts | 153 ++++++++++++ lib/hooks.ts | 13 + lib/state/state.ts | 12 + lib/state/types.ts | 1 + tests/reconstruct.test.ts | 479 ++++++++++++++++++++++++++++++++++++ 5 files changed, 658 insertions(+) create mode 100644 lib/compress/reconstruct.ts create mode 100644 tests/reconstruct.test.ts diff --git a/lib/compress/reconstruct.ts b/lib/compress/reconstruct.ts new file mode 100644 index 00000000..151f7210 --- /dev/null +++ b/lib/compress/reconstruct.ts @@ -0,0 +1,153 @@ +import type { SessionState, WithParts } from "../state" +import type { Logger } from "../logger" +import { messageHasCompress } from "../messages/query" +import { countTokens } from "../token-utils" +import { buildSearchContext, resolveAnchorMessageId, resolveBoundaryIds, resolveSelection } from "./search" +import { allocateBlockId, allocateRunId, applyCompressionState, wrapCompressedSummary } from "./state" + +interface CompressToolResult { + messageId: string + callId: string | undefined + topic: string + ranges: Array<{ + startId: string + endId: string + summary: string + }> + messageIndex: number +} + +function extractCompressResults(messages: WithParts[]): CompressToolResult[] { + const results: CompressToolResult[] = [] + + for (let i = 0; i < messages.length; i++) { + const message = messages[i] + if (!messageHasCompress(message)) { + continue + } + + const parts = Array.isArray(message.parts) ? message.parts : [] + for (const part of parts) { + if (part.type !== "tool" || part.tool !== "compress" || part.state?.status !== "completed") { + continue + } + + const input = part.state?.input + if (!input || !Array.isArray(input.content)) { + continue + } + + const ranges: CompressToolResult["ranges"] = [] + for (const entry of input.content) { + if (typeof entry.startId === "string" && typeof entry.endId === "string" && typeof entry.summary === "string") { + ranges.push({ + startId: entry.startId, + endId: entry.endId, + summary: entry.summary, + }) + } + } + + if (ranges.length === 0) { + continue + } + + results.push({ + messageId: message.info.id, + callId: typeof part.callID === "string" ? part.callID : undefined, + topic: typeof input.topic === "string" ? input.topic : "", + ranges, + messageIndex: i, + }) + } + } + + return results +} + +export function hasCompressHistory(messages: WithParts[]): boolean { + for (const message of messages) { + if (messageHasCompress(message)) { + return true + } + } + return false +} + +export function reconstructFromHistory( + state: SessionState, + logger: Logger, + messages: WithParts[], +): number { + const compressResults = extractCompressResults(messages) + if (compressResults.length === 0) { + return 0 + } + + let totalReconstructed = 0 + let skippedResults = 0 + + for (const result of compressResults) { + const searchContext = buildSearchContext(state, messages) + const runId = allocateRunId(state) + + for (const range of result.ranges) { + try { + const { startReference, endReference } = resolveBoundaryIds( + searchContext, + state, + range.startId, + range.endId, + ) + + const selection = resolveSelection(searchContext, startReference, endReference) + const anchorMessageId = resolveAnchorMessageId(startReference) + + const blockId = allocateBlockId(state) + const storedSummary = wrapCompressedSummary(blockId, range.summary) + const summaryTokens = countTokens(storedSummary) + + const consumedBlockIds = selection.requiredBlockIds + + applyCompressionState( + state, + { + topic: result.topic, + batchTopic: result.topic, + startId: range.startId, + endId: range.endId, + mode: "range", + runId, + compressMessageId: result.messageId, + compressCallId: result.callId, + summaryTokens, + }, + selection, + anchorMessageId, + blockId, + storedSummary, + consumedBlockIds, + ) + + totalReconstructed++ + } catch (err: any) { + skippedResults++ + logger.warn("Skipped reconstruction of compress range", { + startId: range.startId, + endId: range.endId, + error: err.message, + }) + } + } + } + + if (totalReconstructed > 0 || skippedResults > 0) { + logger.info("Reconstructed compression state from history", { + reconstructed: totalReconstructed, + skipped: skippedResults, + totalCompressResults: compressResults.length, + }) + } + + return totalReconstructed +} diff --git a/lib/hooks.ts b/lib/hooks.ts index 67030f1c..945a5ea5 100644 --- a/lib/hooks.ts +++ b/lib/hooks.ts @@ -37,6 +37,7 @@ import { import { type HostPermissionSnapshot } from "./host-permissions" import { compressPermission, syncCompressPermissionState } from "./compress-permission" import { checkSession, ensureSessionInitialized, saveSessionState, syncToolCache } from "./state" +import { reconstructFromHistory } from "./compress/reconstruct" import { cacheSystemPromptTokens } from "./ui/utils" const INTERNAL_AGENT_SIGNATURES = [ @@ -125,6 +126,18 @@ export function createChatMessageTransformHandler( stripHallucinations(output.messages) cacheSystemPromptTokens(state, output.messages) assignMessageRefs(state, output.messages) + + if (state.needsReconstruction) { + const reconstructed = reconstructFromHistory(state, logger, output.messages) + state.needsReconstruction = false + if (reconstructed > 0) { + logger.info("Compression state reconstructed from conversation history", { + blocks: reconstructed, + }) + await saveSessionState(state, logger) + } + } + syncCompressionBlocks(state, logger, output.messages) syncToolCache(state, config, logger, output.messages) buildToolIdList(state, output.messages) diff --git a/lib/state/state.ts b/lib/state/state.ts index 2a8cff6a..485f38d3 100644 --- a/lib/state/state.ts +++ b/lib/state/state.ts @@ -1,6 +1,7 @@ import type { SessionState, ToolParameterEntry, WithParts } from "./types" import type { Logger } from "../logger" import { applyPendingCompressionDurations } from "../compress/timing" +import { hasCompressHistory } from "../compress/reconstruct" import { loadManualModeSetting, loadSessionState, saveSessionState } from "./persistence" import { isSubAgentSession, @@ -52,6 +53,11 @@ export const checkSession = async ( timestamp: lastCompactionTimestamp, }) + if (hasCompressHistory(messages)) { + state.needsReconstruction = true + logger.info("Post-compaction compression history detected — reconstruction needed") + } + saveSessionState(state, logger).catch((error) => { logger.warn("Failed to persist state reset after compaction", { error: error instanceof Error ? error.message : String(error), @@ -99,6 +105,7 @@ export function createSessionState(): SessionState { currentTurn: 0, modelContextLimit: undefined, systemPromptTokens: undefined, + needsReconstruction: false, } } @@ -133,6 +140,7 @@ export function resetSessionState(state: SessionState): void { state.currentTurn = 0 state.modelContextLimit = undefined state.systemPromptTokens = undefined + state.needsReconstruction = false } export async function ensureSessionInitialized( @@ -164,6 +172,10 @@ export async function ensureSessionInitialized( const persisted = await loadSessionState(sessionId, logger) if (persisted === null) { + if (hasCompressHistory(messages)) { + state.needsReconstruction = true + logger.info("New session with compression history detected — reconstruction needed") + } return } diff --git a/lib/state/types.ts b/lib/state/types.ts index acce05f1..9a7f2832 100644 --- a/lib/state/types.ts +++ b/lib/state/types.ts @@ -108,4 +108,5 @@ export interface SessionState { currentTurn: number modelContextLimit: number | undefined systemPromptTokens: number | undefined + needsReconstruction: boolean } diff --git a/tests/reconstruct.test.ts b/tests/reconstruct.test.ts new file mode 100644 index 00000000..16c5cfda --- /dev/null +++ b/tests/reconstruct.test.ts @@ -0,0 +1,479 @@ +import assert from "node:assert/strict" +import test from "node:test" +import { join } from "node:path" +import { tmpdir } from "node:os" +import { mkdirSync } from "node:fs" +import { createSessionState, type WithParts } from "../lib/state" +import { assignMessageRefs } from "../lib/message-ids" +import { hasCompressHistory, reconstructFromHistory } from "../lib/compress/reconstruct" +import { Logger } from "../lib/logger" + +const testDataHome = join(tmpdir(), `opencode-dcp-reconstruct-tests-${process.pid}`) +const testConfigHome = join(tmpdir(), `opencode-dcp-reconstruct-config-tests-${process.pid}`) + +process.env.XDG_DATA_HOME = testDataHome +process.env.XDG_CONFIG_HOME = testConfigHome + +mkdirSync(testDataHome, { recursive: true }) +mkdirSync(testConfigHome, { recursive: true }) + +function textPart(messageID: string, sessionID: string, id: string, text: string) { + return { + id, + messageID, + sessionID, + type: "text" as const, + text, + } +} + +function compressToolPart( + messageID: string, + sessionID: string, + id: string, + callID: string, + topic: string, + content: Array<{ startId: string; endId: string; summary: string }>, +) { + return { + id, + messageID, + sessionID, + callID, + type: "tool" as const, + tool: "compress", + state: { + status: "completed" as const, + input: { topic, content }, + output: `Compressed ${content.length} messages into [Compressed conversation section].`, + }, + } +} + +function buildForkedSessionMessages(sessionID: string): WithParts[] { + return [ + { + info: { + id: "fork-msg-user-1", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 100 }, + } as WithParts["info"], + parts: [textPart("fork-msg-user-1", sessionID, "p1", "Hello")], + }, + { + info: { + id: "fork-msg-assistant-1", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 200 }, + } as WithParts["info"], + parts: [textPart("fork-msg-assistant-1", sessionID, "p2", "Hi there, how can I help?")], + }, + { + info: { + id: "fork-msg-user-2", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 300 }, + } as WithParts["info"], + parts: [textPart("fork-msg-user-2", sessionID, "p3", "Tell me about the system")], + }, + { + info: { + id: "fork-msg-assistant-2", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 400 }, + } as WithParts["info"], + parts: [ + textPart("fork-msg-assistant-2", sessionID, "p4", "The system has three components"), + compressToolPart( + "fork-msg-assistant-2", + sessionID, + "p5", + "call-compress-1", + "Initial greeting", + [ + { + startId: "m0001", + endId: "m0002", + summary: "User greeted and assistant responded with offer to help.", + }, + ], + ), + ], + }, + { + info: { + id: "fork-msg-user-3", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 500 }, + } as WithParts["info"], + parts: [textPart("fork-msg-user-3", sessionID, "p6", "Can you compress more?")], + }, + ] +} + +test("hasCompressHistory returns true when messages contain completed compress tool calls", () => { + const messages = buildForkedSessionMessages("ses_test_1") + assert.equal(hasCompressHistory(messages), true) +}) + +test("hasCompressHistory returns false when no compress tool calls exist", () => { + const sessionID = "ses_test_2" + const messages: WithParts[] = [ + { + info: { + id: "msg-user-1", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 100 }, + } as WithParts["info"], + parts: [textPart("msg-user-1", sessionID, "p1", "Hello")], + }, + { + info: { + id: "msg-assistant-1", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 200 }, + } as WithParts["info"], + parts: [textPart("msg-assistant-1", sessionID, "p2", "Hi there")], + }, + ] + assert.equal(hasCompressHistory(messages), false) +}) + +test("reconstructFromHistory rebuilds compression state from compress tool results", () => { + const sessionID = "ses_fork_reconstruct_1" + const messages = buildForkedSessionMessages(sessionID) + const state = createSessionState() + state.sessionId = sessionID + const logger = new Logger(false) + + assignMessageRefs(state, messages) + + assert.equal(state.messageIds.byRef.get("m0001"), "fork-msg-user-1") + assert.equal(state.messageIds.byRef.get("m0002"), "fork-msg-assistant-1") + + const reconstructed = reconstructFromHistory(state, logger, messages) + + assert.equal(reconstructed, 1) + assert.equal(state.prune.messages.blocksById.size, 1) + assert.equal(state.prune.messages.activeBlockIds.size, 1) + + const block = state.prune.messages.blocksById.get(1) + assert.ok(block) + assert.equal(block.active, true) + assert.equal(block.topic, "Initial greeting") + assert.equal(block.compressMessageId, "fork-msg-assistant-2") + assert.equal(block.compressCallId, "call-compress-1") + assert.equal(block.startId, "m0001") + assert.equal(block.endId, "m0002") + + const entry1 = state.prune.messages.byMessageId.get("fork-msg-user-1") + assert.ok(entry1) + assert.ok(entry1.activeBlockIds.includes(1)) + + const entry2 = state.prune.messages.byMessageId.get("fork-msg-assistant-1") + assert.ok(entry2) + assert.ok(entry2.activeBlockIds.includes(1)) +}) + +test("reconstructFromHistory handles multiple sequential compressions", () => { + const sessionID = "ses_fork_reconstruct_2" + const messages: WithParts[] = [ + { + info: { + id: "msg-u1", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 100 }, + } as WithParts["info"], + parts: [textPart("msg-u1", sessionID, "p1", "First message")], + }, + { + info: { + id: "msg-a1", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 200 }, + } as WithParts["info"], + parts: [textPart("msg-a1", sessionID, "p2", "First response")], + }, + { + info: { + id: "msg-u2", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 300 }, + } as WithParts["info"], + parts: [textPart("msg-u2", sessionID, "p3", "Second message")], + }, + { + info: { + id: "msg-a2", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 400 }, + } as WithParts["info"], + parts: [ + textPart("msg-a2", sessionID, "p4", "Second response"), + compressToolPart("msg-a2", sessionID, "p5", "call-1", "First compression", [ + { + startId: "m0001", + endId: "m0002", + summary: "First exchange summary.", + }, + ]), + ], + }, + { + info: { + id: "msg-u3", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 500 }, + } as WithParts["info"], + parts: [textPart("msg-u3", sessionID, "p6", "Third message")], + }, + { + info: { + id: "msg-a3", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 600 }, + } as WithParts["info"], + parts: [ + textPart("msg-a3", sessionID, "p7", "Third response"), + compressToolPart("msg-a3", sessionID, "p8", "call-2", "Second compression", [ + { + startId: "m0003", + endId: "m0004", + summary: "Second exchange summary.", + }, + ]), + ], + }, + { + info: { + id: "msg-u4", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 700 }, + } as WithParts["info"], + parts: [textPart("msg-u4", sessionID, "p9", "Fourth message")], + }, + ] + + const state = createSessionState() + state.sessionId = sessionID + const logger = new Logger(false) + + assignMessageRefs(state, messages) + const reconstructed = reconstructFromHistory(state, logger, messages) + + assert.equal(reconstructed, 2) + assert.equal(state.prune.messages.blocksById.size, 2) + + const block1 = state.prune.messages.blocksById.get(1) + assert.ok(block1) + assert.equal(block1.topic, "First compression") + assert.equal(block1.compressMessageId, "msg-a2") + + const block2 = state.prune.messages.blocksById.get(2) + assert.ok(block2) + assert.equal(block2.topic, "Second compression") + assert.equal(block2.compressMessageId, "msg-a3") +}) + +test("reconstructFromHistory handles block refs (bN) in ranges by consuming earlier blocks", () => { + const sessionID = "ses_fork_reconstruct_3" + const messages: WithParts[] = [ + { + info: { + id: "msg-u1", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 100 }, + } as WithParts["info"], + parts: [textPart("msg-u1", sessionID, "p1", "First message")], + }, + { + info: { + id: "msg-a1", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 200 }, + } as WithParts["info"], + parts: [textPart("msg-a1", sessionID, "p2", "First response")], + }, + { + info: { + id: "msg-u2", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 300 }, + } as WithParts["info"], + parts: [textPart("msg-u2", sessionID, "p3", "Second message")], + }, + { + info: { + id: "msg-a2", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 400 }, + } as WithParts["info"], + parts: [ + textPart("msg-a2", sessionID, "p4", "Second response with compress"), + compressToolPart("msg-a2", sessionID, "p5", "call-1", "First block", [ + { + startId: "m0001", + endId: "m0002", + summary: "First exchange compressed.", + }, + ]), + ], + }, + { + info: { + id: "msg-u3", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 500 }, + } as WithParts["info"], + parts: [textPart("msg-u3", sessionID, "p6", "Third message")], + }, + { + info: { + id: "msg-a3", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 600 }, + } as WithParts["info"], + parts: [textPart("msg-a3", sessionID, "p7", "Third response")], + }, + { + info: { + id: "msg-u4", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 700 }, + } as WithParts["info"], + parts: [textPart("msg-u4", sessionID, "p8", "Fourth message")], + }, + { + info: { + id: "msg-a4", + role: "assistant", + sessionID, + agent: "assistant", + time: { created: 800 }, + } as WithParts["info"], + parts: [ + textPart("msg-a4", sessionID, "p9", "Fourth response with mega compress"), + compressToolPart("msg-a4", sessionID, "p10", "call-2", "Mega block", [ + { + startId: "b1", + endId: "m0006", + summary: "Everything up to third response compressed.", + }, + ]), + ], + }, + { + info: { + id: "msg-u5", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 900 }, + } as WithParts["info"], + parts: [textPart("msg-u5", sessionID, "p11", "Latest message")], + }, + ] + + const state = createSessionState() + state.sessionId = sessionID + const logger = new Logger(false) + + assignMessageRefs(state, messages) + const reconstructed = reconstructFromHistory(state, logger, messages) + + assert.equal(reconstructed, 2) + + const block1 = state.prune.messages.blocksById.get(1) + assert.ok(block1) + assert.equal(block1.active, false) + + const block2 = state.prune.messages.blocksById.get(2) + assert.ok(block2) + assert.equal(block2.active, true) + assert.equal(block2.startId, "b1") + assert.equal(block2.endId, "m0006") + + assert.equal(state.prune.messages.activeBlockIds.size, 1) + assert.ok(state.prune.messages.activeBlockIds.has(2)) +}) + +test("reconstructFromHistory returns 0 when no compress results exist", () => { + const sessionID = "ses_fork_no_compress" + const messages: WithParts[] = [ + { + info: { + id: "msg-u1", + role: "user", + sessionID, + agent: "assistant", + model: { providerID: "anthropic", modelID: "claude-test" }, + time: { created: 100 }, + } as WithParts["info"], + parts: [textPart("msg-u1", sessionID, "p1", "Hello")], + }, + ] + + const state = createSessionState() + state.sessionId = sessionID + const logger = new Logger(false) + + assignMessageRefs(state, messages) + const reconstructed = reconstructFromHistory(state, logger, messages) + assert.equal(reconstructed, 0) + assert.equal(state.prune.messages.blocksById.size, 0) +})