fix: reconstruct compression state after fork and compaction

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
+}

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)

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
     }
 

lib/state/types.ts

@@ -108,4 +108,5 @@ export interface SessionState {
     currentTurn: number
     modelContextLimit: number | undefined
     systemPromptTokens: number | undefined
+    needsReconstruction: boolean
 }