diff --git a/src/preflight.ts b/src/preflight.ts index b901f4f6..67050a99 100644 --- a/src/preflight.ts +++ b/src/preflight.ts @@ -1,5 +1,6 @@ import { defaultCountTokens, + collectBlockContent, viableRanges, type CompressionCore, type Config, @@ -167,23 +168,21 @@ function rangeChars(messages: CoreMessage[], startIdx: number, endIdx: number): return chars; } -function renderRange(messages: CoreMessage[], startIdx: number, endIdx: number): string { - const parts: string[] = []; - for (let i = startIdx; i <= endIdx && i < messages.length; i++) { - const m = messages[i]; - const text = (m.text ?? "").trim(); - if (!text) continue; - const label = - m.contentType === "tool-call" - ? `assistant tool-call ${m.toolName ?? "?"}` - : m.contentType === "tool-result" - ? `tool result ${m.toolName ?? "?"}` - : m.contentType === "reasoning" - ? "assistant reasoning" - : m.role; - parts.push(`[${label}]\n${text}`); +function splitSummaryContent(content: string, budget: number, countTokens: (text: string) => number): string[] { + const chunks: string[] = []; + let offset = 0; + while (offset < content.length) { + let low = offset + 1; + let high = content.length; + while (low < high) { + const mid = Math.ceil((low + high) / 2); + if (countTokens(content.slice(offset, mid)) <= budget) low = mid; + else high = mid - 1; + } + chunks.push(content.slice(offset, low)); + offset = low; } - return parts.join("\n\n"); + return chunks; } // minUnits: never close a chunk below this many countText units while more @@ -474,7 +473,7 @@ export async function preflightCompress(deps: PreflightDeps, messages: CoreMessa // estimate (not currentTokens, which is floored by a possibly-stale // lastInputTokens from a prior model): if the real payload already // fits, stop instead of folding protected content. - if (!relaxed && result.payloadEstimate >= limit) { + if (!relaxed && (baselineKnown ? result.payloadEstimate : finalUpper) >= limit) { activeConfig = relaxedConfig(deps.config); relaxed = true; // #575-merge: the summarization budget counts per protection @@ -522,17 +521,36 @@ export async function preflightCompress(deps: PreflightDeps, messages: CoreMessa const startRef = maps.idxToRef.get(cs); const endRef = maps.idxToRef.get(ce); if (!startRef || !endRef) continue; - if (rangeChars(messages, cs, ce) < minChars) continue; - const content = renderRange(messages, cs, ce); + const preview = deps.core.applyCompression({ + messages, + state: deps.session.state, + config: activeConfig, + ranges: [{ startRef, endRef, summary: "x".repeat(Math.max(MIN_SUMMARY_CHARS, activeConfig.compress.minSummaryLength)) }], + }); + const previousBlockIds = new Set(deps.session.state.blocks.map((block) => block.blockId)); + const planned = preview.state.blocks.find((block) => !previousBlockIds.has(block.blockId)); + if (!planned) continue; + // Use the original state so consumed child blocks are still active and render as summaries. + const content = collectBlockContent(deps.session.state, planned, messages, { full: false }).text; if (content.length === 0) continue; - if (summaryCalls >= MAX_SUMMARY_CALLS_PER_PREFLIGHT) { - budgetHit = true; - break; - } - summaryCalls += 1; - let summary: string | null; + let summary: string | null = null; try { - summary = await summarizeRange(deps, content, startRef, endRef); + const parts: string[] = []; + const chunks = splitSummaryContent(content, budget, countText); + for (const chunk of chunks) { + if (summaryCalls >= MAX_SUMMARY_CALLS_PER_PREFLIGHT) { + budgetHit = true; + break; + } + summaryCalls += 1; + const part = await summarizeRange(deps, chunk, startRef, endRef); + if (!part) break; + parts.push(part); + } + if (!budgetHit && parts.length === chunks.length) { + const candidate = parts.join("\n\n"); + if (activeConfig.compress.maxSummaryLength <= 0 || candidate.length <= activeConfig.compress.maxSummaryLength) summary = candidate || null; + } } catch (err) { if (err instanceof UpstreamHttpError) { failure = { @@ -576,7 +594,7 @@ export async function preflightCompress(deps: PreflightDeps, messages: CoreMessa // baseline currentTokens is char-based, so net the folded span's // char count against it instead of the token-based credit. const compressed = deps.session.stats.compressCreditTokens - creditBefore; - const folded = baselineKnown ? compressed : rangeChars(messages, cs, ce); + const folded = baselineKnown ? compressed : messages.filter((message) => planned.directMessageIds.includes(message.id)).reduce((total, message) => total + (message.text ?? "").length, 0); currentTokens = Math.max(0, currentTokens - folded + countText(summary)); deps.session.stats.lastInputTokens += defaultCountTokens(summary); appliedThisRound += 1; @@ -586,7 +604,17 @@ export async function preflightCompress(deps: PreflightDeps, messages: CoreMessa if (appliedThisRound > 0) break; if (failure || budgetHit) break; } - if (appliedThisRound === 0) break; + if (appliedThisRound === 0) { + if (!failure && !budgetHit && !relaxed && (baselineKnown ? result.payloadEstimate : finalUpper) >= limit) { + activeConfig = relaxedConfig(deps.config); + relaxed = true; + summaryCalls = 0; + budgetHit = false; + deps.log("warn", "[preflight] no usable ranges outside the protected recent zone; relaxing soft protection (preserveRecentMessages/Tokens -> 0) and retrying"); + continue; + } + break; + } } if (currentTokens >= limit && !failure) { if (budgetHit) { diff --git a/tests/preflight-covered-ranges.test.ts b/tests/preflight-covered-ranges.test.ts new file mode 100644 index 00000000..676fd816 --- /dev/null +++ b/tests/preflight-covered-ranges.test.ts @@ -0,0 +1,111 @@ +import assert from "node:assert/strict"; +import http from "node:http"; +import { once } from "node:events"; +import test from "node:test"; +import { createCore, createInitialState, defaultConfig, defaultPrompts, type CoreMessage } from "acp-kernel"; +import { preflightCompress } from "../src/preflight.ts"; +import type { Session } from "../src/session.ts"; + +process.env.NODE_ENV = "test"; +const SUMMARY = "The preceding work is summarized here with all decisions and remaining tasks preserved for the next turn. No historical tool calls should be repeated."; + +async function runCoveredRange(protectLatest = false, pairBoundary = false, unknownBaseline = false, segmentFailure?: "missing" | "over-limit", maxSummaryLength?: number) { + const core = createCore(); + const config = defaultConfig(unknownBaseline ? 100000 : 272000); + config.preserveRecentMessages = 5; + config.preserveRecentTokens = 5000; + config.compress.minCompressRange = 5000; + if (segmentFailure === "over-limit") config.compress.maxSummaryLength = 200; + if (maxSummaryLength !== undefined) config.compress.maxSummaryLength = maxSummaryLength; + config.protectedTools = protectLatest ? ["large_result"] : []; + const messages: CoreMessage[] = [ + { id: "early", role: "user", contentType: "text", text: "EARLY ".repeat(pairBoundary ? 1000 : 100) }, + { id: "covered", role: "assistant", contentType: "text", text: "HIDDEN_RAW ".repeat(80000) }, + ...Array.from({ length: 10 }, (_, i): CoreMessage => ({ id: `middle-${i}`, role: i % 2 ? "assistant" : "user", contentType: "text", text: `MIDDLE_${i} `.repeat(500) })), + { id: "large-call", role: "assistant", contentType: "tool-call", toolName: "large_result", toolCallId: "large", text: "{}" }, + { id: "large-result", role: "tool", contentType: "tool-result", toolName: "large_result", toolCallId: "large", text: "LATEST_LARGE ".repeat(unknownBaseline ? 10000 : 90000) }, + ]; + const turn = core.processTurn({ messages, state: createInitialState(), config: { ...config, modelContextLimit: 27200000 }, tokenCount: 300000, renderTags: "text-only" }); + const coveredRef = turn.state.messageRefs.byRaw.covered; + const compressed = core.applyCompression({ messages, state: turn.state, config: { ...config, preserveRecentMessages: 0, preserveRecentTokens: 0 }, ranges: [{ startRef: coveredRef, endRef: coveredRef, summary: SUMMARY }] }); + assert.equal(compressed.result.blocksCreated, 1); + const session: Session = { + id: "covered-preflight", meta: {}, metadata: {}, state: compressed.state, + stats: { requests: 1, tokensSaved: 0, inputTokens: 0, cachedTokens: 0, outputTokens: 0, cacheSamples: 0, lastInputTokens: unknownBaseline ? 0 : 287565, compressCreditTokens: 0, contextTokens: 287565 }, + createdAt: Date.now(), lastSeen: Date.now(), blockContents: new Map(), inFlight: 0, persisted: false, + }; + const summaries: string[] = []; + let largeSegments = 0; + const upstream = http.createServer((req, res) => { + const chunks: Buffer[] = []; + req.on("data", (chunk: Buffer) => chunks.push(chunk)); + req.on("end", () => { + const body = JSON.parse(Buffer.concat(chunks).toString("utf8")) as { input: { content: string }[] }; + const content = body.input[0].content; + summaries.push(content); + res.setHeader("content-type", "application/json"); + if (content.includes("LATEST_LARGE")) largeSegments += 1; + res.end(JSON.stringify({ output_text: segmentFailure === "missing" && content.includes("LATEST_LARGE") && largeSegments % 2 === 0 ? "" : SUMMARY })); + }); + }); + upstream.listen(0, "127.0.0.1"); + await once(upstream, "listening"); + const logs: string[] = []; + try { + const result = await preflightCompress({ core, session, config, prompts: defaultPrompts, protocol: "responses", url: `http://127.0.0.1:${(upstream.address() as { port: number }).port}`, headers: {}, model: "test", unknownBaseline, log: (_level, message) => logs.push(message) }, messages); + return { result, summaries, logs, session }; + } finally { + const closed = once(upstream, "close"); + upstream.close(); + upstream.closeAllConnections(); + await closed; + } +} + +test("preflight skips fully covered raw chunks and relaxes unusable normal ranges to fold a large recent result", async () => { + const { result, summaries, logs } = await runCoveredRange(); + assert.equal(result.fitsWindow, true, JSON.stringify(result)); + assert.ok(result.compressedRanges > 0); + assert.ok(summaries.join("").includes("LATEST_LARGE ".repeat(90000)), "the complete tool result, including the pair-expanded tail, must reach the summarizer"); + assert.ok(summaries.every((content) => content.length <= Math.floor(272000 * 0.6) * 4), "each summary input obeys the chunk budget"); + assert.ok(summaries.every((content) => !content.includes("HIDDEN_RAW")), "covered raw chunks must not be resummarized by raw ref"); + assert.ok(logs.some((message) => message.includes("relaxing soft protection"))); + assert.ok(logs.every((message) => !message.includes("already covered"))); +}); + +test("preflight keeps hard-protected tools excluded after unusable normal ranges", async () => { + const { result, summaries } = await runCoveredRange(true); + assert.equal(result.fitsWindow, false); + assert.ok(summaries.every((content) => !content.includes("LATEST_LARGE"))); + assert.ok(summaries.length <= 16, "both protection regimes have bounded summary calls"); +}); + + +test("preflight summarizes the complete tool pair before the kernel consumes its expanded result boundary", async () => { + const { result, summaries } = await runCoveredRange(false, true); + assert.equal(result.fitsWindow, true); + assert.ok(summaries.join("").includes("LATEST_LARGE ".repeat(90000)), "no tool-result tail may be folded without entering a summary request"); +}); + +test("unknown-baseline preflight relaxes protection using the conservative upper bound", async () => { + const { result, summaries } = await runCoveredRange(false, false, true); + assert.equal(result.fitsWindow, true, JSON.stringify(result)); + assert.ok(summaries.join("").includes("LATEST_LARGE ".repeat(10000))); +}); + + +for (const failure of ["missing", "over-limit"] as const) { + test(`preflight does not apply an incomplete segmented summary (${failure})`, async () => { + const { result, session } = await runCoveredRange(false, true, false, failure); + assert.equal(result.fitsWindow, false); + assert.ok(session.state.blocks.filter((block) => block.active).every((block) => !block.effectiveMessageIds.includes("large-result")), "the original tool result remains available when any summary part is missing or too long"); + }); +} + +for (const limit of [0, -1]) { + test(`preflight accepts unlimited summary length (${limit})`, async () => { + const { result } = await runCoveredRange(false, true, false, undefined, limit); + assert.equal(result.fitsWindow, true); + assert.ok(result.compressedRanges > 0); + }); +}