fix(compaction): clamp summary output tokens

Fixes #4390.
This commit is contained in:
Armin Ronacher
2026-05-11 16:36:27 +02:00
parent f348a06294
commit 3d9e14d748
4 changed files with 70 additions and 9 deletions

View File

@@ -538,7 +538,10 @@ export async function generateSummary(
previousSummary?: string,
thinkingLevel?: ThinkingLevel,
): Promise<string> {
const maxTokens = Math.floor(0.8 * reserveTokens);
const maxTokens = Math.min(
Math.floor(0.8 * reserveTokens),
model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY,
);
// Use update prompt if we have a previous summary, otherwise initial prompt
let basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT;
@@ -808,7 +811,10 @@ async function generateTurnPrefixSummary(
signal?: AbortSignal,
thinkingLevel?: ThinkingLevel,
): Promise<string> {
const maxTokens = Math.floor(0.5 * reserveTokens); // Smaller budget for turn prefix
const maxTokens = Math.min(
Math.floor(0.5 * reserveTokens),
model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY,
); // Smaller budget for turn prefix
const llmMessages = convertToLlm(messages);
const conversationText = serializeConversation(llmMessages);
const promptText = `<conversation>\n${conversationText}\n</conversation>\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`;

View File

@@ -1,7 +1,7 @@
import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { AssistantMessage, Model } from "@earendil-works/pi-ai";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { generateSummary } from "../src/core/compaction/index.js";
import { type CompactionPreparation, compact, generateSummary } from "../src/core/compaction/index.js";
const { completeSimpleMock } = vi.hoisted(() => ({
completeSimpleMock: vi.fn(),
@@ -15,7 +15,7 @@ vi.mock("@earendil-works/pi-ai", async (importOriginal) => {
};
});
function createModel(reasoning: boolean): Model<"anthropic-messages"> {
function createModel(reasoning: boolean, maxTokens = 8192): Model<"anthropic-messages"> {
return {
id: reasoning ? "reasoning-model" : "non-reasoning-model",
name: reasoning ? "Reasoning Model" : "Non-reasoning Model",
@@ -26,7 +26,7 @@ function createModel(reasoning: boolean): Model<"anthropic-messages"> {
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 200000,
maxTokens: 8192,
maxTokens,
};
}
@@ -115,4 +115,20 @@ describe("generateSummary reasoning options", () => {
});
expect(completeSimpleMock.mock.calls[0][2]).not.toHaveProperty("reasoning");
});
it("clamps compaction summary maxTokens to the model output cap", async () => {
const preparation: CompactionPreparation = {
firstKeptEntryId: "entry-keep",
messagesToSummarize: messages,
turnPrefixMessages: messages,
isSplitTurn: true,
tokensBefore: 600000,
fileOps: { read: new Set(), written: new Set(), edited: new Set() },
settings: { enabled: true, reserveTokens: 500000, keepRecentTokens: 20000 },
};
await compact(preparation, createModel(false, 128000), "test-key");
expect(completeSimpleMock.mock.calls.map((call) => call[2]?.maxTokens)).toEqual([128000, 128000]);
});
});