test(agents): share context pruning trim fixtures

This commit is contained in:
Vincent Koc
2026-04-12 05:19:41 +01:00
parent 0acfa47e08
commit aa415b2553

View File

@@ -10,6 +10,9 @@ type AssistantContentBlock = AssistantMessage["content"][number];
const CONTEXT_WINDOW_1M = {
model: { contextWindow: 1_000_000 },
} as unknown as ExtensionContext;
const CONTEXT_WINDOW_5K = {
model: { contextWindow: 5_000 },
} as unknown as ExtensionContext;
function makeUser(text: string): AgentMessage {
return {
@@ -58,6 +61,43 @@ function makeToolResult(
} as AgentMessage;
}
function pruneWithOversizedAssistantThinking(params: {
assistantBlock: AssistantContentBlock;
dropThinkingBlocksForEstimate?: boolean;
}) {
return pruneContextMessages({
messages: [
makeUser("hello"),
makeToolResult([{ type: "text", text: "X".repeat(2_000) }]),
makeAssistant([params.assistantBlock, { type: "text", text: "done" }]),
],
settings: {
...buildToolTrimSettings(),
},
ctx: CONTEXT_WINDOW_5K,
isToolPrunable: () => true,
...(params.dropThinkingBlocksForEstimate ? { dropThinkingBlocksForEstimate: true } : {}),
});
}
function buildToolTrimSettings() {
return {
keepLastAssistants: 1,
softTrimRatio: 0.5,
softTrim: { maxChars: 200, headChars: 100, tailChars: 50 },
hardClear: { ...DEFAULT_CONTEXT_PRUNING_SETTINGS.hardClear, enabled: false },
};
}
function expectToolResultWasTrimmed(result: AgentMessage[]) {
const toolResult = result.find((message) => message.role === "toolResult") as Extract<
AgentMessage,
{ role: "toolResult" }
>;
const textBlock = toolResult.content[0] as { type: "text"; text: string };
expect(textBlock.text).toContain("[Tool result trimmed:");
}
describe("pruneContextMessages", () => {
it("does not crash on assistant message with malformed thinking block (missing thinking string)", () => {
const messages: AgentMessage[] = [
@@ -124,74 +164,26 @@ describe("pruneContextMessages", () => {
});
it("counts thinkingSignature bytes when estimating assistant message size", () => {
const messages: AgentMessage[] = [
makeUser("hello"),
makeToolResult([{ type: "text", text: "X".repeat(2_000) }]),
makeAssistant([
{
type: "thinking",
thinking: "[redacted]",
thinkingSignature: "S".repeat(40_000),
redacted: true,
} as unknown as AssistantContentBlock,
{ type: "text", text: "done" },
]),
];
const result = pruneContextMessages({
messages,
settings: {
...DEFAULT_CONTEXT_PRUNING_SETTINGS,
keepLastAssistants: 1,
softTrimRatio: 0.5,
softTrim: { maxChars: 200, headChars: 100, tailChars: 50 },
hardClear: { ...DEFAULT_CONTEXT_PRUNING_SETTINGS.hardClear, enabled: false },
},
ctx: { model: { contextWindow: 5_000 } } as unknown as ExtensionContext,
isToolPrunable: () => true,
const result = pruneWithOversizedAssistantThinking({
assistantBlock: {
type: "thinking",
thinking: "[redacted]",
thinkingSignature: "S".repeat(40_000),
redacted: true,
} as unknown as AssistantContentBlock,
});
const toolResult = result.find((message) => message.role === "toolResult") as Extract<
AgentMessage,
{ role: "toolResult" }
>;
const textBlock = toolResult.content[0] as { type: "text"; text: string };
expect(textBlock.text).toContain("[Tool result trimmed:");
expectToolResultWasTrimmed(result);
});
it("counts redacted_thinking data bytes when estimating assistant message size", () => {
const messages: AgentMessage[] = [
makeUser("hello"),
makeToolResult([{ type: "text", text: "X".repeat(2_000) }]),
makeAssistant([
{
type: "redacted_thinking",
data: "D".repeat(40_000),
thinkingSignature: "sig",
} as unknown as AssistantContentBlock,
{ type: "text", text: "done" },
]),
];
const result = pruneContextMessages({
messages,
settings: {
...DEFAULT_CONTEXT_PRUNING_SETTINGS,
keepLastAssistants: 1,
softTrimRatio: 0.5,
softTrim: { maxChars: 200, headChars: 100, tailChars: 50 },
hardClear: { ...DEFAULT_CONTEXT_PRUNING_SETTINGS.hardClear, enabled: false },
},
ctx: { model: { contextWindow: 5_000 } } as unknown as ExtensionContext,
isToolPrunable: () => true,
const result = pruneWithOversizedAssistantThinking({
assistantBlock: {
type: "redacted_thinking",
data: "D".repeat(40_000),
thinkingSignature: "sig",
} as unknown as AssistantContentBlock,
});
const toolResult = result.find((message) => message.role === "toolResult") as Extract<
AgentMessage,
{ role: "toolResult" }
>;
const textBlock = toolResult.content[0] as { type: "text"; text: string };
expect(textBlock.text).toContain("[Tool result trimmed:");
expectToolResultWasTrimmed(result);
});
it("ignores non-latest thinking signatures that will be dropped before send", () => {
@@ -214,12 +206,9 @@ describe("pruneContextMessages", () => {
messages,
settings: {
...DEFAULT_CONTEXT_PRUNING_SETTINGS,
keepLastAssistants: 1,
softTrimRatio: 0.5,
softTrim: { maxChars: 200, headChars: 100, tailChars: 50 },
hardClear: { ...DEFAULT_CONTEXT_PRUNING_SETTINGS.hardClear, enabled: false },
...buildToolTrimSettings(),
},
ctx: { model: { contextWindow: 5_000 } } as unknown as ExtensionContext,
ctx: CONTEXT_WINDOW_5K,
isToolPrunable: () => true,
dropThinkingBlocksForEstimate: true,
});