From 95f75302e4afa276638ceeac64f13e38f58a796d Mon Sep 17 00:00:00 2001 From: ranxianglei Date: Sat, 22 Aug 2026 22:17:26 +0800 Subject: [PATCH] fix: match nested blocks by any visible representation in range (billion-context-pi#195 follow-up) --- src/boundaries.ts | 29 +++++- src/compress.ts | 7 +- src/index.ts | 1 + tests/regression-promote-after-prune.test.ts | 103 +++++++++++++++++++ 4 files changed, 135 insertions(+), 5 deletions(-) diff --git a/src/boundaries.ts b/src/boundaries.ts index 2fd9680..1854c68 100644 --- a/src/boundaries.ts +++ b/src/boundaries.ts @@ -131,8 +131,7 @@ export function resolveBoundaries( const nestedBlockIds: string[] = []; const nestedSeen = new Set(); for (const block of activeBlocks(input.state)) { - const anchor = visibleBlockAnchor(block, indexByMessageId); - if (anchor !== null && anchor >= startIndex && anchor <= endIndex) { + if (blockVisibleInRange(block, indexByMessageId, startIndex, endIndex)) { if (!nestedSeen.has(block.blockId)) { nestedSeen.add(block.blockId); nestedBlockIds.push(block.blockId); @@ -311,6 +310,32 @@ export function visibleBlockAnchor( return earliestIndexOfIds(block.effectiveMessageIds, indexByMessageId); } +// A block participates in a range when ANY of its visible representations +// (rendered summary or earliest surviving raw) falls inside it. The summary +// alone is not sufficient: when a block covers the session's first user +// message, prune keeps that one raw and inserts the summary BEFORE it, so the +// summary index can sit outside a range that still contains the raw. +export function blockVisibleInRange( + block: CompressionBlock, + indexByMessageId: Map, + startIndex: number, + endIndex: number, +): boolean { + const summaryIndex = indexByMessageId.get(summaryMessageId(block.blockId)); + if ( + summaryIndex !== undefined && + summaryIndex >= startIndex && + summaryIndex <= endIndex + ) { + return true; + } + const rawIndex = earliestIndexOfIds( + block.effectiveMessageIds, + indexByMessageId, + ); + return rawIndex !== null && rawIndex >= startIndex && rawIndex <= endIndex; +} + export function earliestIndexOfIds( ids: string[], indexByMessageId: Map, diff --git a/src/compress.ts b/src/compress.ts index fe25cda..b7de336 100644 --- a/src/compress.ts +++ b/src/compress.ts @@ -8,7 +8,7 @@ import { validateConfig } from "./config.js"; import { BoundaryNotFoundError, resolveBoundaries, - visibleBlockAnchor, + blockVisibleInRange, } from "./boundaries.js"; import type { ResolvedRange } from "./boundaries.js"; import { truncateLargeToolOutputs } from "./truncate-tools.js"; @@ -666,8 +666,9 @@ function applySingleRange(input: SingleRangeInput): SingleRangeOutcome { const nestedSeen = new Set(resolved.nestedBlockIds); for (const block of activeBlocks(input.state)) { if (nestedSeen.has(block.blockId)) continue; - const anchor = visibleBlockAnchor(block, indexByMessageId); - if (anchor !== null && anchor >= adjustedStart && anchor <= adjustedEnd) { + if ( + blockVisibleInRange(block, indexByMessageId, adjustedStart, adjustedEnd) + ) { nestedSeen.add(block.blockId); resolved.nestedBlockIds.push(block.blockId); } diff --git a/src/index.ts b/src/index.ts index 09848ce..a67b5a8 100644 --- a/src/index.ts +++ b/src/index.ts @@ -40,6 +40,7 @@ export { parseBoundary, BoundaryNotFoundError, visibleBlockAnchor, + blockVisibleInRange, } from "./boundaries.js"; export { defaultCountTokens, diff --git a/tests/regression-promote-after-prune.test.ts b/tests/regression-promote-after-prune.test.ts index 267809d..6bbb14c 100644 --- a/tests/regression-promote-after-prune.test.ts +++ b/tests/regression-promote-after-prune.test.ts @@ -545,3 +545,106 @@ test("promote-after-prune: compressing a consumed block's bN ref snaps to its ac false, ); }); + +test("promote-after-prune: m-ref range consumes a block that covers the first user message (summary anchor precedes the surviving raw)", () => { + const core = createCore(); + const cfg = config({ preserveRecentMessages: 0 }); + const messages = [ + msg("raw-1", "u1 ".repeat(200)), + msg("raw-2", "a2 ".repeat(200), "assistant"), + msg("raw-3", "u3 ".repeat(200)), + msg("raw-4", "a4 ".repeat(200), "assistant"), + msg("raw-5", "u5 ".repeat(200)), + msg("raw-6", "a6 ".repeat(200), "assistant"), + msg("raw-7", "u7 ".repeat(200)), + msg("raw-8", "a8 ".repeat(200), "assistant"), + msg("raw-9", "u9 ".repeat(200)), + msg("raw-10", "a10 ".repeat(200), "assistant"), + msg("raw-11", "u11 ".repeat(200)), + msg("raw-12", "a12 ".repeat(200), "assistant"), + msg("raw-13", "u13 ".repeat(200)), + ]; + const state = makeState( + [ + { + blockId: "b1", + effectiveMessageIds: ["raw-1", "raw-2", "raw-3", "raw-4", "raw-5"], + }, + ], + 2, + ); + + const turn = core.processTurn({ + messages, + state, + config: cfg, + tokenCount: 5000, + }); + const visibleIds = turn.messages.map((m) => m.id); + // Prune keeps the first user message and inserts the summary at the block's + // earliest raw position — BEFORE the surviving raw, so the summary index + // differs from the raw's index. + assert.ok(visibleIds.includes(summaryMessageId("b1"))); + assert.ok(visibleIds.includes("raw-1"), "first user message survives prune"); + assert.ok(!visibleIds.includes("raw-5"), "covered raws are hidden"); + assert.ok( + visibleIds.indexOf(summaryMessageId("b1")) < visibleIds.indexOf("raw-1"), + "summary precedes the surviving raw", + ); + + const startRef = turn.state.messageRefs.byRaw["raw-1"]!; + const endRef = turn.state.messageRefs.byRaw["raw-13"]!; + + const applied = core.applyCompression({ + ranges: [{ startRef, endRef, summary: "S".repeat(80) }], + messages: turn.messages, + state: turn.state, + config: cfg, + }); + assert.deepEqual( + applied.result.errors, + [], + `unexpected errors: ${applied.result.errors.join("; ")}`, + ); + assert.equal(applied.result.blocksCreated, 1); + + const newBlock = applied.state.blocks[applied.state.blocks.length - 1]!; + // b1 must be consumed via its surviving raw even though its summary anchor + // sits outside the resolved range. + assert.deepEqual(newBlock.directBlockIds, ["b1"]); + assert.deepEqual(newBlock.effectiveMessageIds.sort(), [ + "raw-1", + "raw-10", + "raw-11", + "raw-12", + "raw-13", + "raw-2", + "raw-3", + "raw-4", + "raw-5", + "raw-6", + "raw-7", + "raw-8", + "raw-9", + ]); + assert.deepEqual(newBlock.directMessageIds.sort(), [ + "raw-10", + "raw-11", + "raw-12", + "raw-13", + "raw-6", + "raw-7", + "raw-8", + "raw-9", + ]); + for (const id of [ + ...newBlock.effectiveMessageIds, + ...newBlock.directMessageIds, + ]) { + assert.ok(!isSummaryMessageId(id), `synthetic id leaked into block: ${id}`); + } + assert.equal( + applied.state.blocks.find((b) => b.blockId === "b1")!.active, + false, + ); +});