mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-12 10:12:11 +03:00
fix: bound quadratic session-dedup memory growth (#7855)
* fix: bound long-context compression memory * perf: scan session dedup line starts natively --------- Co-authored-by: Ravi Tharuma <RaviTharuma@users.noreply.github.com>
This commit is contained in:
@@ -46,6 +46,15 @@ const ENGINE_ID = "session-dedup";
|
||||
const DEFAULT_MIN_BLOCK_CHARS = 80;
|
||||
/** Minimum number of lines a block must span to be a dedup candidate. */
|
||||
const MIN_BLOCK_LINES = 3;
|
||||
/**
|
||||
* Request-wide ceiling for the suffix strings materialized by the exact pass.
|
||||
* 32 MiB keeps ordinary sessions byte-identical while preventing line-rich inputs
|
||||
* from retaining a quadratic graph of suffix copies.
|
||||
*/
|
||||
const MAX_SUFFIX_WORK_CHARS = 32 * 1024 * 1024;
|
||||
const SUFFIX_WORK_BUDGET_WARNING = "session-dedup: skipped (suffix work budget exceeded)";
|
||||
|
||||
type SuffixWorkBudget = { remaining: number };
|
||||
|
||||
// ─── hash helper (SHA-256 prefix, collision-resistant) ───────────────────────
|
||||
|
||||
@@ -58,6 +67,24 @@ function hashBlock(text: string): string {
|
||||
|
||||
// ─── suffix-block extraction ──────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Reserves the characters that findSuffixBlocks() would materialize for one text.
|
||||
* The scan observes line starts without splitting or constructing any suffix strings.
|
||||
*/
|
||||
function reserveSuffixWork(text: string, passCount: number, budget: SuffixWorkBudget): boolean {
|
||||
let start = 0;
|
||||
while (start <= text.length) {
|
||||
const suffixChars = (text.length - start) * passCount;
|
||||
if (suffixChars > budget.remaining) return false;
|
||||
budget.remaining -= suffixChars;
|
||||
|
||||
const nextNewline = text.indexOf("\n", start);
|
||||
if (nextNewline === -1) break;
|
||||
start = nextNewline + 1;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* For each starting line position, emit the suffix block `lines[start..end]`
|
||||
* (i.e. from `start` to the end of the line array). This ensures that any
|
||||
@@ -118,7 +145,9 @@ function dedupeWithinMessage(
|
||||
|
||||
for (const { block } of sortedBlocks) {
|
||||
// Only dedup blocks that appear 2+ times in the text.
|
||||
const occurrences = (result.match(new RegExp(block.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"), "g")) || []).length;
|
||||
const occurrences = (
|
||||
result.match(new RegExp(block.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"), "g")) || []
|
||||
).length;
|
||||
if (occurrences < 2) continue;
|
||||
|
||||
const sha = hashBlock(block);
|
||||
@@ -240,7 +269,7 @@ type MessageLike = {
|
||||
function processMessages(
|
||||
messages: MessageLike[],
|
||||
minBlockChars: number
|
||||
): { messages: MessageLike[]; dedupCount: number } {
|
||||
): { messages: MessageLike[]; dedupCount: number; suffixWorkBudgetExceeded: boolean } {
|
||||
// Collect (msgIdx, text) for non-system string-content messages.
|
||||
// For multipart, index each text part separately.
|
||||
const msgTexts: Array<{ msgIdx: number; text: string }> = [];
|
||||
@@ -262,13 +291,24 @@ function processMessages(
|
||||
}
|
||||
|
||||
if (msgTexts.length === 0) {
|
||||
return { messages, dedupCount: 0 };
|
||||
return { messages, dedupCount: 0, suffixWorkBudgetExceeded: false };
|
||||
}
|
||||
|
||||
// Single-message exact dedup enumerates suffixes once; cross-message dedup does so
|
||||
// in both passes. Reserve the request-wide work up front so no quadratic suffix graph
|
||||
// is partially materialized before the engine decides to fail open.
|
||||
const suffixWorkBudget: SuffixWorkBudget = { remaining: MAX_SUFFIX_WORK_CHARS };
|
||||
const passCount = msgTexts.length === 1 ? 1 : 2;
|
||||
for (const { text } of msgTexts) {
|
||||
if (!reserveSuffixWork(text, passCount, suffixWorkBudget)) {
|
||||
return { messages, dedupCount: 0, suffixWorkBudgetExceeded: true };
|
||||
}
|
||||
}
|
||||
|
||||
const { deduped, dedupCount } = dedupMessageTexts(msgTexts, minBlockChars);
|
||||
|
||||
if (dedupCount === 0) {
|
||||
return { messages, dedupCount: 0 };
|
||||
return { messages, dedupCount: 0, suffixWorkBudgetExceeded: false };
|
||||
}
|
||||
|
||||
const result = messages.map((msg, i) => {
|
||||
@@ -297,7 +337,7 @@ function processMessages(
|
||||
return { ...msg };
|
||||
});
|
||||
|
||||
return { messages: result, dedupCount };
|
||||
return { messages: result, dedupCount, suffixWorkBudgetExceeded: false };
|
||||
}
|
||||
|
||||
// ─── schema & validation ──────────────────────────────────────────────────────
|
||||
@@ -343,7 +383,8 @@ function validateSessionDedupConfig(config: Record<string, unknown>): EngineVali
|
||||
const f = config["fuzzy"];
|
||||
if (typeof f === "object" && f !== null) {
|
||||
const fe = (f as Record<string, unknown>)["enabled"];
|
||||
if (fe !== undefined && typeof fe !== "boolean") errors.push("fuzzy.enabled must be a boolean");
|
||||
if (fe !== undefined && typeof fe !== "boolean")
|
||||
errors.push("fuzzy.enabled must be a boolean");
|
||||
} else if (typeof f !== "boolean") {
|
||||
errors.push("fuzzy must be an object { enabled } or a boolean");
|
||||
}
|
||||
@@ -394,10 +435,18 @@ export const sessionDedupEngine: CompressionEngine = {
|
||||
}
|
||||
|
||||
const start = performance.now();
|
||||
const { messages: exactMessages, dedupCount } = processMessages(
|
||||
messages as MessageLike[],
|
||||
minBlockChars
|
||||
);
|
||||
const {
|
||||
messages: exactMessages,
|
||||
dedupCount,
|
||||
suffixWorkBudgetExceeded,
|
||||
} = processMessages(messages as MessageLike[], minBlockChars);
|
||||
|
||||
if (suffixWorkBudgetExceeded) {
|
||||
const durationMs = Math.round(performance.now() - start);
|
||||
const stats = createCompressionStats(body, body, "stacked", [], undefined, durationMs);
|
||||
stats.validationWarnings = [SUFFIX_WORK_BUDGET_WARNING];
|
||||
return { body, compressed: false, stats };
|
||||
}
|
||||
|
||||
const { messages: finalMessages, fuzzyCount } = runFuzzyPass(
|
||||
exactMessages,
|
||||
|
||||
32
tests/fixtures/compression/session-dedup-memory-7849.ts
vendored
Normal file
32
tests/fixtures/compression/session-dedup-memory-7849.ts
vendored
Normal file
@@ -0,0 +1,32 @@
|
||||
import { applyStackedCompression } from "../../../open-sse/services/compression/index.ts";
|
||||
|
||||
const lines = Array.from({ length: 4_000 }, (_, index) => {
|
||||
const prefix = `line-${index.toString().padStart(4, "0")}:`;
|
||||
return prefix + "x".repeat(80 - prefix.length);
|
||||
});
|
||||
const body = {
|
||||
messages: [{ role: "tool", content: lines.join("\n") }],
|
||||
};
|
||||
const enginesRun: string[] = [];
|
||||
const result = applyStackedCompression(
|
||||
body,
|
||||
[
|
||||
{ engine: "session-dedup" },
|
||||
{ engine: "lite" },
|
||||
{ engine: "rtk" },
|
||||
{ engine: "headroom" },
|
||||
{ engine: "caveman" },
|
||||
],
|
||||
{
|
||||
onEngineStep: (step) => {
|
||||
enginesRun.push(step.engine);
|
||||
},
|
||||
}
|
||||
);
|
||||
|
||||
process.stdout.write(
|
||||
JSON.stringify({
|
||||
enginesRun,
|
||||
warnings: result.stats?.validationWarnings ?? [],
|
||||
})
|
||||
);
|
||||
138
tests/unit/compression/session-dedup-memory-7849.test.ts
Normal file
138
tests/unit/compression/session-dedup-memory-7849.test.ts
Normal file
@@ -0,0 +1,138 @@
|
||||
import assert from "node:assert/strict";
|
||||
import { spawnSync } from "node:child_process";
|
||||
import { dirname, join } from "node:path";
|
||||
import test from "node:test";
|
||||
import { fileURLToPath } from "node:url";
|
||||
|
||||
import { sessionDedupEngine } from "../../../open-sse/services/compression/engines/session-dedup/index.ts";
|
||||
|
||||
const REPO_ROOT = join(dirname(fileURLToPath(import.meta.url)), "../../..");
|
||||
const FIXTURE = join(REPO_ROOT, "tests/fixtures/compression/session-dedup-memory-7849.ts");
|
||||
const SUFFIX_WORK_BUDGET = 32 * 1024 * 1024;
|
||||
const SUFFIX_WORK_BUDGET_WARNING = "session-dedup: skipped (suffix work budget exceeded)";
|
||||
|
||||
function makeFixedWidthText(lineCount: number, lineChars: number, tag: string): string {
|
||||
return Array.from({ length: lineCount }, (_, index) => {
|
||||
const prefix = `${tag}-${index.toString().padStart(4, "0")}:`;
|
||||
assert.ok(prefix.length <= lineChars);
|
||||
return prefix + "x".repeat(lineChars - prefix.length);
|
||||
}).join("\n");
|
||||
}
|
||||
|
||||
function projectedSuffixWork(text: string, passCount: number): number {
|
||||
let work = 0;
|
||||
for (let start = 0; start <= text.length; start++) {
|
||||
if (start === 0 || text.charCodeAt(start - 1) === 10) {
|
||||
work += (text.length - start) * passCount;
|
||||
}
|
||||
}
|
||||
return work;
|
||||
}
|
||||
|
||||
function makeSharedBudgetBody(): Record<string, unknown> {
|
||||
return {
|
||||
messages: [
|
||||
{ role: "tool", content: makeFixedWidthText(600, 49, "first") },
|
||||
{ role: "tool", content: makeFixedWidthText(600, 49, "second") },
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
test("#7849: shares the two-pass suffix-work budget across all messages", () => {
|
||||
const body = makeSharedBudgetBody();
|
||||
const messages = body.messages as Array<{ content: string }>;
|
||||
const perMessageWork = messages.map(({ content }) => projectedSuffixWork(content, 2));
|
||||
|
||||
assert.ok(
|
||||
perMessageWork.every((work) => work < SUFFIX_WORK_BUDGET),
|
||||
"each message must fit the two-pass budget on its own"
|
||||
);
|
||||
assert.ok(
|
||||
messages.reduce((total, { content }) => total + projectedSuffixWork(content, 1), 0) <
|
||||
SUFFIX_WORK_BUDGET,
|
||||
"the pair must fit if incorrectly charged for only one pass"
|
||||
);
|
||||
assert.ok(
|
||||
perMessageWork.reduce((total, work) => total + work, 0) > SUFFIX_WORK_BUDGET,
|
||||
"the pair must exceed the shared budget when correctly charged for two passes"
|
||||
);
|
||||
|
||||
for (const message of messages) {
|
||||
const individualResult = sessionDedupEngine.apply({
|
||||
messages: [message, { role: "assistant", content: "a unique short companion" }],
|
||||
});
|
||||
assert.equal(individualResult.stats, null, "each message must be accepted individually");
|
||||
}
|
||||
|
||||
const result = sessionDedupEngine.apply(body);
|
||||
assert.deepEqual(result.stats?.validationWarnings, [SUFFIX_WORK_BUDGET_WARNING]);
|
||||
});
|
||||
|
||||
test("#7849: exhausted suffix-work budget fails open with exact zero-savings stats", () => {
|
||||
const body = makeSharedBudgetBody();
|
||||
const result = sessionDedupEngine.apply(body);
|
||||
|
||||
assert.strictEqual(result.body, body, "budget exhaustion must return the input body by identity");
|
||||
assert.equal(result.compressed, false);
|
||||
assert.ok(result.stats, "budget exhaustion must return explanatory stats");
|
||||
assert.equal(result.stats.originalTokens, result.stats.compressedTokens);
|
||||
assert.equal(result.stats.savingsPercent, 0);
|
||||
assert.deepEqual(result.stats.validationWarnings, [SUFFIX_WORK_BUDGET_WARNING]);
|
||||
});
|
||||
|
||||
test("#7849: near-boundary under-budget request still deduplicates", () => {
|
||||
const repeatedText = makeFixedWidthText(578, 49, "same");
|
||||
const projectedWork = projectedSuffixWork(repeatedText, 2) * 2;
|
||||
assert.ok(projectedWork <= SUFFIX_WORK_BUDGET);
|
||||
assert.ok(
|
||||
SUFFIX_WORK_BUDGET - projectedWork < 100_000,
|
||||
"fixture must remain close to the work-budget boundary"
|
||||
);
|
||||
|
||||
const body = {
|
||||
messages: [
|
||||
{ role: "user", content: repeatedText },
|
||||
{ role: "user", content: repeatedText },
|
||||
],
|
||||
};
|
||||
const result = sessionDedupEngine.apply(body);
|
||||
const messages = result.body.messages as Array<{ content: string }>;
|
||||
|
||||
assert.equal(result.compressed, true);
|
||||
assert.equal(messages[0].content, repeatedText);
|
||||
assert.match(messages[1].content, /^\[dedup:ref sha=[0-9a-f]{24}\]$/);
|
||||
assert.ok((result.stats?.savingsPercent ?? 0) > 0);
|
||||
assert.deepEqual(result.stats?.validationWarnings ?? [], []);
|
||||
});
|
||||
|
||||
test(
|
||||
"#7849: line-rich long context stays within a 512 MiB heap and the stacked pipeline continues",
|
||||
{ timeout: 60_000 },
|
||||
() => {
|
||||
const child = spawnSync(
|
||||
process.execPath,
|
||||
["--max-old-space-size=512", "--import", "tsx/esm", FIXTURE],
|
||||
{
|
||||
cwd: REPO_ROOT,
|
||||
encoding: "utf8",
|
||||
timeout: 45_000,
|
||||
}
|
||||
);
|
||||
|
||||
assert.equal(
|
||||
child.status,
|
||||
0,
|
||||
`compression child must not OOM or time out\nstdout: ${child.stdout}\nstderr: ${child.stderr}`
|
||||
);
|
||||
|
||||
const output = JSON.parse(child.stdout) as {
|
||||
enginesRun: string[];
|
||||
warnings: string[];
|
||||
};
|
||||
assert.deepEqual(output.enginesRun, ["session-dedup", "lite", "rtk", "headroom", "caveman"]);
|
||||
assert.ok(
|
||||
output.warnings.includes("session-dedup: skipped (suffix work budget exceeded)"),
|
||||
`expected an explicit session-dedup work-budget warning, got ${JSON.stringify(output.warnings)}`
|
||||
);
|
||||
}
|
||||
);
|
||||
Reference in New Issue
Block a user