@@ -0,0 +1,154 @@
1+import type { AgentMessage } from "@mariozechner/pi-agent-core";
2+import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
3+import { beforeEach, describe, expect, it, vi } from "vitest";
4+5+const piCodingAgentMocks = vi.hoisted(() => ({
6+estimateTokens: vi.fn((message: unknown) => Math.ceil(JSON.stringify(message).length / 4)),
7+generateSummary: vi.fn(),
8+}));
9+10+vi.mock("@mariozechner/pi-coding-agent", async () => {
11+const actual = await vi.importActual<typeof import("@mariozechner/pi-coding-agent")>(
12+"@mariozechner/pi-coding-agent",
13+);
14+return {
15+ ...actual,
16+estimateTokens: piCodingAgentMocks.estimateTokens,
17+generateSummary: piCodingAgentMocks.generateSummary,
18+};
19+});
20+21+const mockGenerateSummary = piCodingAgentMocks.generateSummary;
22+23+let summarizeInStages: typeof import("./compaction.js").summarizeInStages;
24+25+async function loadFreshCompactionModuleForTest() {
26+vi.resetModules();
27+({ summarizeInStages } = await import("./compaction.js"));
28+}
29+30+function makeMessage(index: number, size = 1200): AgentMessage {
31+return {
32+role: "user",
33+content: `m${index}-${"x".repeat(size)}`,
34+timestamp: index,
35+};
36+}
37+38+describe("compaction reserveTokens clamping", () => {
39+beforeEach(async () => {
40+await loadFreshCompactionModuleForTest();
41+mockGenerateSummary.mockReset();
42+mockGenerateSummary.mockResolvedValue("summary");
43+piCodingAgentMocks.estimateTokens.mockReset();
44+piCodingAgentMocks.estimateTokens.mockImplementation((message: unknown) =>
45+Math.ceil(JSON.stringify(message).length / 4),
46+);
47+});
48+49+it("clamps reserveTokens when model maxTokens is smaller than requested", async () => {
50+// Simulate the exact bug scenario: large context window (1M) with
51+// reserveTokensFloor of 300K, but model output limit is only 128K.
52+// Without clamping, generateSummary would receive 300K and compute
53+// max_tokens = floor(0.8 * 300K) = 240K, exceeding the 128K model limit.
54+const model = {
55+provider: "anthropic",
56+model: "claude-sonnet-4-6",
57+contextWindow: 1_000_000,
58+maxTokens: 128_000,
59+} as unknown as NonNullable<ExtensionContext["model"]>;
60+61+await summarizeInStages({
62+ model,
63+apiKey: "test-key", // pragma: allowlist secret
64+reserveTokens: 300_000,
65+maxChunkTokens: 8000,
66+contextWindow: 1_000_000,
67+signal: new AbortController().signal,
68+messages: [makeMessage(1), makeMessage(2)],
69+});
70+71+expect(mockGenerateSummary).toHaveBeenCalled();
72+// Third argument to generateSummary is reserveTokens.
73+// With maxTokens 128K, the clamp should be floor(128_000 / 0.8) = 160_000.
74+const passedReserveTokens = mockGenerateSummary.mock.calls[0][2];
75+expect(passedReserveTokens).toBeLessThanOrEqual(Math.floor(128_000 / 0.8));
76+expect(passedReserveTokens).toBe(160_000);
77+});
78+79+it("does not clamp when model maxTokens is large enough", async () => {
80+const model = {
81+provider: "anthropic",
82+model: "claude-opus-4-6",
83+contextWindow: 200_000,
84+maxTokens: 32_000,
85+} as unknown as NonNullable<ExtensionContext["model"]>;
86+87+// reserveTokens 4000 is well under floor(32_000 / 0.8) = 40_000
88+await summarizeInStages({
89+ model,
90+apiKey: "test-key", // pragma: allowlist secret
91+reserveTokens: 4000,
92+maxChunkTokens: 8000,
93+contextWindow: 200_000,
94+signal: new AbortController().signal,
95+messages: [makeMessage(1), makeMessage(2)],
96+});
97+98+expect(mockGenerateSummary).toHaveBeenCalled();
99+const passedReserveTokens = mockGenerateSummary.mock.calls[0][2];
100+expect(passedReserveTokens).toBe(4000);
101+});
102+103+it("falls back to 128K default when model has no maxTokens field", async () => {
104+// Model without maxTokens defined — should default to 128_000 as the cap.
105+const model = {
106+provider: "anthropic",
107+model: "claude-3-opus",
108+contextWindow: 1_000_000,
109+} as unknown as NonNullable<ExtensionContext["model"]>;
110+111+await summarizeInStages({
112+ model,
113+apiKey: "test-key", // pragma: allowlist secret
114+reserveTokens: 300_000,
115+maxChunkTokens: 8000,
116+contextWindow: 1_000_000,
117+signal: new AbortController().signal,
118+messages: [makeMessage(1), makeMessage(2)],
119+});
120+121+expect(mockGenerateSummary).toHaveBeenCalled();
122+// Fallback maxTokens is 128_000, so clamp = floor(128_000 / 0.8) = 160_000
123+const passedReserveTokens = mockGenerateSummary.mock.calls[0][2];
124+expect(passedReserveTokens).toBe(160_000);
125+});
126+127+it("clamps consistently across all chunks in staged summarization", async () => {
128+const model = {
129+provider: "anthropic",
130+model: "claude-sonnet-4-6",
131+contextWindow: 1_000_000,
132+maxTokens: 128_000,
133+} as unknown as NonNullable<ExtensionContext["model"]>;
134+135+// Use enough messages and small chunk size to force multiple chunks
136+await summarizeInStages({
137+ model,
138+apiKey: "test-key", // pragma: allowlist secret
139+reserveTokens: 300_000,
140+maxChunkTokens: 1000,
141+contextWindow: 1_000_000,
142+signal: new AbortController().signal,
143+messages: Array.from({ length: 4 }, (_, i) => makeMessage(i + 1)),
144+parts: 2,
145+minMessagesForSplit: 4,
146+});
147+148+expect(mockGenerateSummary.mock.calls.length).toBeGreaterThan(1);
149+const expectedClamp = Math.floor(128_000 / 0.8);
150+for (const call of mockGenerateSummary.mock.calls) {
151+expect(call[2]).toBeLessThanOrEqual(expectedClamp);
152+}
153+});
154+});
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。