| import type { AgentMessage, AgentTool } from "@earendil-works/pi-agent-core"; |
| import { |
| type AssistantMessage, |
| type Context, |
| createAssistantMessageEventStream, |
| fauxAssistantMessage, |
| fauxToolCall, |
| type Model, |
| type SimpleStreamOptions, |
| } from "@earendil-works/pi-ai"; |
| import { Type } from "typebox"; |
| import { afterEach, describe, expect, it, vi } from "vitest"; |
| import { estimateTokens } from "../../src/core/compaction/index.ts"; |
| import { createHarness, getUserTexts, type Harness } from "./harness.ts"; |
|
|
| type SessionWithCompactionInternals = { |
| _checkCompaction: (assistantMessage: AssistantMessage, skipAbortedCheck?: boolean) => Promise<boolean>; |
| _runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise<boolean>; |
| }; |
|
|
| function createUsage(totalTokens: number) { |
| return { |
| input: totalTokens, |
| output: 0, |
| cacheRead: 0, |
| cacheWrite: 0, |
| totalTokens, |
| cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, |
| }; |
| } |
|
|
| function createAssistant( |
| harness: Harness, |
| options: { |
| stopReason?: AssistantMessage["stopReason"]; |
| errorMessage?: string; |
| totalTokens?: number; |
| timestamp?: number; |
| }, |
| ): AssistantMessage { |
| const model = harness.getModel(); |
| return { |
| ...fauxAssistantMessage("", { |
| stopReason: options.stopReason, |
| errorMessage: options.errorMessage, |
| timestamp: options.timestamp, |
| }), |
| api: model.api, |
| provider: model.provider, |
| model: model.id, |
| usage: createUsage(options.totalTokens ?? 0), |
| }; |
| } |
|
|
| function useSummaryStreamFn( |
| harness: Harness, |
| summary: string, |
| onRequest?: (context: Context, options: SimpleStreamOptions | undefined) => void, |
| ): () => number { |
| let callCount = 0; |
| harness.session.agent.streamFunction = (model, context, options) => { |
| callCount++; |
| onRequest?.(context, options); |
| const stream = createAssistantMessageEventStream(); |
| queueMicrotask(() => { |
| const message: AssistantMessage = { |
| ...fauxAssistantMessage(summary), |
| api: model.api, |
| provider: model.provider, |
| model: model.id, |
| usage: createUsage(10), |
| }; |
| stream.push({ type: "done", reason: "stop", message }); |
| }); |
| return stream; |
| }; |
| return () => callCount; |
| } |
|
|
| function seedCompactableSession(harness: Harness): void { |
| harness.settingsManager.applyOverrides({ compaction: { keepRecentTokens: 1 } }); |
| const now = Date.now(); |
| harness.sessionManager.appendMessage({ |
| role: "user", |
| content: [{ type: "text", text: "message to compact" }], |
| timestamp: now - 1000, |
| }); |
| const assistant = createAssistant(harness, { |
| stopReason: "stop", |
| totalTokens: 100, |
| timestamp: now - 500, |
| }); |
| assistant.content = [{ type: "text", text: "assistant response to compact" }]; |
| harness.sessionManager.appendMessage(assistant); |
| harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; |
| } |
|
|
| async function createAbortableCompactionHarness(): Promise<{ |
| harness: Harness; |
| compactionStarted: Promise<void>; |
| }> { |
| let markCompactionStarted = () => {}; |
| const compactionStarted = new Promise<void>((resolve) => { |
| markCompactionStarted = resolve; |
| }); |
| const harness = await createHarness({ |
| settings: { compaction: { keepRecentTokens: 1 } }, |
| extensionFactories: [ |
| (pi) => { |
| pi.on("session_before_compact", async (event) => { |
| return await new Promise<{ cancel: true }>((resolve) => { |
| event.signal.addEventListener("abort", () => resolve({ cancel: true }), { once: true }); |
| markCompactionStarted(); |
| }); |
| }); |
| }, |
| ], |
| }); |
| seedCompactableSession(harness); |
| return { harness, compactionStarted }; |
| } |
|
|
| describe("AgentSession compaction characterization", () => { |
| const harnesses: Harness[] = []; |
|
|
| afterEach(() => { |
| vi.useRealTimers(); |
| vi.restoreAllMocks(); |
| while (harnesses.length > 0) { |
| harnesses.pop()?.cleanup(); |
| } |
| }); |
|
|
| it("manually compacts using an extension-provided summary", async () => { |
| const summaryUsage = { |
| input: 10, |
| output: 20, |
| cacheRead: 30, |
| cacheWrite: 40, |
| totalTokens: 100, |
| cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, |
| }; |
| const harness = await createHarness({ |
| settings: { compaction: { keepRecentTokens: 1 } }, |
| extensionFactories: [ |
| (pi) => { |
| pi.on("session_before_compact", async (event) => ({ |
| compaction: { |
| summary: "summary from extension", |
| firstKeptEntryId: event.preparation.firstKeptEntryId, |
| tokensBefore: event.preparation.tokensBefore, |
| usage: summaryUsage, |
| details: { source: "extension" }, |
| }, |
| })); |
| }, |
| ], |
| }); |
| harnesses.push(harness); |
|
|
| await harness.session.prompt("one"); |
| await harness.session.prompt("two"); |
| const statsBefore = harness.session.getSessionStats(); |
|
|
| const result = await harness.session.compact(); |
| const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); |
| const estimatedTokensAfter = harness.session.messages.reduce((sum, message) => sum + estimateTokens(message), 0); |
|
|
| expect(result.summary).toBe("summary from extension"); |
| expect(result.usage).toEqual(summaryUsage); |
| expect(result.estimatedTokensAfter).toBe(estimatedTokensAfter); |
| expect(compactionEntries).toHaveLength(1); |
| const compactionEntry = compactionEntries[0]; |
| if (compactionEntry?.type === "compaction") { |
| expect(compactionEntry.usage).toEqual(summaryUsage); |
| } |
| const statsAfter = harness.session.getSessionStats(); |
| expect(statsAfter.tokens.input).toBe(statsBefore.tokens.input + summaryUsage.input); |
| expect(statsAfter.tokens.output).toBe(statsBefore.tokens.output + summaryUsage.output); |
| expect(statsAfter.tokens.cacheRead).toBe(statsBefore.tokens.cacheRead + summaryUsage.cacheRead); |
| expect(statsAfter.tokens.cacheWrite).toBe(statsBefore.tokens.cacheWrite + summaryUsage.cacheWrite); |
| expect(statsAfter.cost).toBe(statsBefore.cost + summaryUsage.cost.total); |
| expect(harness.session.messages[0]?.role).toBe("compactionSummary"); |
| }); |
|
|
| it("allows a queued prompt to start when manual compaction ends", async () => { |
| const harness = await createHarness({ |
| settings: { compaction: { keepRecentTokens: 1 } }, |
| extensionFactories: [ |
| (pi) => { |
| pi.on("session_before_compact", async (event) => ({ |
| compaction: { |
| summary: "manual compacted", |
| firstKeptEntryId: event.preparation.firstKeptEntryId, |
| tokensBefore: event.preparation.tokensBefore, |
| details: {}, |
| }, |
| })); |
| }, |
| ], |
| }); |
| harnesses.push(harness); |
| seedCompactableSession(harness); |
| harness.setResponses([fauxAssistantMessage("queued response")]); |
|
|
| let queuedPrompt: Promise<void> | undefined; |
| harness.session.subscribe((event) => { |
| if (event.type === "compaction_end" && event.reason === "manual" && event.result) { |
| expect(harness.session.isCompacting).toBe(false); |
| queuedPrompt = harness.session.prompt("queued after compaction"); |
| } |
| }); |
|
|
| await harness.session.compact(); |
| if (!queuedPrompt) throw new Error("compaction_end did not start the queued prompt"); |
| await queuedPrompt; |
|
|
| expect(getUserTexts(harness)).toContain("queued after compaction"); |
| expect(harness.session.getLastAssistantText()).toBe("queued response"); |
| }); |
|
|
| it("throws when compacting without a model", async () => { |
| const harness = await createHarness(); |
| harnesses.push(harness); |
| harness.session.agent.state.model = undefined as unknown as Model<any>; |
|
|
| await expect(harness.session.compact()).rejects.toThrow("No model selected"); |
| }); |
|
|
| it("throws when compacting without configured auth", async () => { |
| const harness = await createHarness({ withConfiguredAuth: false }); |
| harnesses.push(harness); |
|
|
| await expect(harness.session.compact()).rejects.toThrow(`No API key found for ${harness.getModel().provider}.`); |
| }); |
|
|
| it("manually compacts with a custom streamFn when registry auth is absent", async () => { |
| const harness = await createHarness({ withConfiguredAuth: false }); |
| harnesses.push(harness); |
| seedCompactableSession(harness); |
| const getStreamCallCount = useSummaryStreamFn(harness, "summary from custom stream"); |
|
|
| const result = await harness.session.compact(); |
|
|
| expect(result.summary).toContain("summary from custom stream"); |
| expect(getStreamCallCount()).toBe(1); |
| }); |
|
|
| it("manually compacts with provider-resolved bearer auth", async () => { |
| const harness = await createHarness({ withConfiguredAuth: false }); |
| harnesses.push(harness); |
| const model = harness.getModel(); |
| harness.session.modelRuntime.registerNativeProvider({ |
| id: model.provider, |
| name: "Faux bearer provider", |
| auth: { |
| apiKey: { |
| name: "Faux bearer token", |
| resolve: async () => ({ |
| auth: { headers: { Authorization: "Bearer ambient-token" } }, |
| source: "ambient bearer token", |
| }), |
| }, |
| }, |
| getModels: () => harness.models, |
| stream: () => createAssistantMessageEventStream(), |
| streamSimple: () => createAssistantMessageEventStream(), |
| }); |
| seedCompactableSession(harness); |
| harness.setResponses([ |
| (_context, options) => { |
| expect(options?.apiKey).toBeUndefined(); |
| expect(options?.headers).toEqual({ Authorization: "Bearer ambient-token" }); |
| return fauxAssistantMessage("summary with bearer auth"); |
| }, |
| ]); |
|
|
| const result = await harness.session.compact(); |
|
|
| expect(result.summary).toContain("summary with bearer auth"); |
| expect(harness.faux.state.callCount).toBe(1); |
| }); |
|
|
| it("uses the standalone compaction request context", async () => { |
| const harness = await createHarness({ settings: { compaction: { keepRecentTokens: 1 } } }); |
| harnesses.push(harness); |
| seedCompactableSession(harness); |
|
|
| const transformContext = vi.fn(async (messages: AgentMessage[]) => messages); |
| harness.session.agent.transformContext = transformContext; |
| harness.session.agent.sessionId = "active-routing-session"; |
| harness.session.agent.transport = "websocket"; |
|
|
| let requestContext: Context | undefined; |
| let requestOptions: SimpleStreamOptions | undefined; |
| useSummaryStreamFn(harness, "standalone summary", (context, options) => { |
| requestContext = context; |
| requestOptions = options; |
| }); |
|
|
| await harness.session.compact(); |
|
|
| expect(transformContext).not.toHaveBeenCalled(); |
| expect(requestContext?.systemPrompt).not.toBe(harness.session.agent.state.systemPrompt); |
| expect(requestContext?.tools).toBeUndefined(); |
| expect(JSON.stringify(requestContext?.messages)).toContain("<conversation>"); |
| expect(requestOptions).toMatchObject({ cacheRetention: "none" }); |
| expect(requestOptions?.sessionId).not.toBe("active-routing-session"); |
| expect(requestOptions?.transport).toBeUndefined(); |
| }); |
|
|
| it("persists usage from pi-generated manual compaction", async () => { |
| const harness = await createHarness({ withConfiguredAuth: false }); |
| harnesses.push(harness); |
| seedCompactableSession(harness); |
| useSummaryStreamFn(harness, "summary from custom stream"); |
|
|
| const result = await harness.session.compact(); |
|
|
| const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); |
| expect(result.usage).toEqual(createUsage(10)); |
| expect(compactionEntries).toHaveLength(1); |
| expect(compactionEntries[0]?.type === "compaction" ? compactionEntries[0].usage : undefined).toEqual( |
| createUsage(10), |
| ); |
| }); |
|
|
| it("auto-compacts with a custom streamFn when registry auth is absent", async () => { |
| const harness = await createHarness({ withConfiguredAuth: false }); |
| harnesses.push(harness); |
| seedCompactableSession(harness); |
| const getStreamCallCount = useSummaryStreamFn(harness, "auto summary from custom stream"); |
| const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; |
|
|
| await sessionInternals._runAutoCompaction("threshold", false); |
|
|
| const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); |
| const compactionEnd = harness.eventsOfType("compaction_end").at(-1); |
| expect(compactionEntries).toHaveLength(1); |
| expect(compactionEnd?.result?.estimatedTokensAfter).toBeGreaterThan(0); |
| expect(getStreamCallCount()).toBe(1); |
| }); |
|
|
| it("notifies extensions when auto-compaction fails", async () => { |
| const failedEvents: Array<{ |
| reason: "manual" | "threshold" | "overflow"; |
| errorMessage?: string; |
| aborted: boolean; |
| willRetry: boolean; |
| fromExtension: boolean; |
| }> = []; |
| const harness = await createHarness({ |
| extensionFactories: [ |
| (pi) => { |
| pi.on("session_compact_failed", async (event) => { |
| failedEvents.push(event); |
| }); |
| }, |
| ], |
| }); |
| harnesses.push(harness); |
| seedCompactableSession(harness); |
| harness.session.agent.streamFunction = () => { |
| throw new Error("summary generator blew up"); |
| }; |
| const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; |
|
|
| await expect(sessionInternals._runAutoCompaction("threshold", false)).resolves.toBe(false); |
|
|
| expect(harness.eventsOfType("compaction_end").at(-1)).toMatchObject({ |
| reason: "threshold", |
| aborted: false, |
| willRetry: false, |
| errorMessage: "Auto-compaction failed: summary generator blew up", |
| }); |
| expect(failedEvents).toEqual([ |
| expect.objectContaining({ |
| type: "session_compact_failed", |
| reason: "threshold", |
| aborted: false, |
| willRetry: false, |
| fromExtension: false, |
| errorMessage: "Auto-compaction failed: summary generator blew up", |
| }), |
| ]); |
| }); |
|
|
| it("compacts and resumes after a length stop below the desired output limit", async () => { |
| const harness = await createHarness({ |
| models: [{ id: "faux-1", contextWindow: 1000, maxTokens: 100 }], |
| settings: { compaction: { keepRecentTokens: 1, reserveTokens: 0 } }, |
| extensionFactories: [ |
| (pi) => { |
| pi.on("session_before_compact", async (event) => ({ |
| compaction: { |
| summary: "overflow compacted", |
| firstKeptEntryId: event.preparation.firstKeptEntryId, |
| tokensBefore: event.preparation.tokensBefore, |
| details: {}, |
| }, |
| })); |
| }, |
| ], |
| }); |
| harnesses.push(harness); |
| harness.setResponses([ |
| fauxAssistantMessage("partial response", { stopReason: "length" }), |
| fauxAssistantMessage("completed response"), |
| ]); |
|
|
| await harness.session.prompt("x".repeat(5000)); |
|
|
| expect(harness.faux.state.callCount).toBe(2); |
| expect(harness.eventsOfType("compaction_end").at(-1)).toMatchObject({ |
| reason: "overflow", |
| aborted: false, |
| willRetry: true, |
| }); |
| expect(harness.session.getLastAssistantText()).toBe("completed response"); |
| }); |
|
|
| |
| it.each([false, true])( |
| "compacts after a tool result in the same run (model override: %s)", |
| async (modelOverride) => { |
| const toolResult = `large-tool-result:${"x".repeat(6800)}`; |
| const largeTool: AgentTool = { |
| name: "large_result", |
| label: "Large result", |
| description: "Returns enough content to cross the compaction threshold", |
| parameters: Type.Object({}), |
| execute: async () => ({ content: [{ type: "text", text: toolResult }], details: {} }), |
| }; |
| const order: string[] = []; |
| const observedSettings: unknown[] = []; |
| const harness = await createHarness({ |
| models: [{ id: "faux-1", contextWindow: 2600, maxTokens: 100 }], |
| settings: { |
| compaction: modelOverride |
| ? { |
| enabled: true, |
| reserveTokens: 0, |
| keepRecentTokens: 20000, |
| modelOverrides: { "faux/faux-1": { reserveTokens: 400, keepRecentTokens: 1750 } }, |
| } |
| : { enabled: true, reserveTokens: 400, keepRecentTokens: 1750 }, |
| }, |
| tools: [largeTool], |
| extensionFactories: [ |
| (pi) => { |
| pi.on("session_before_compact", (event) => { |
| order.push("compaction"); |
| observedSettings.push(event.preparation.settings); |
| return { |
| compaction: { |
| summary: "compacted history", |
| firstKeptEntryId: event.preparation.firstKeptEntryId, |
| tokensBefore: event.preparation.tokensBefore, |
| details: {}, |
| }, |
| }; |
| }); |
| }, |
| ], |
| }); |
| harnesses.push(harness); |
| let resumedRequest = ""; |
| harness.setResponses([ |
| fauxAssistantMessage(`old-history:${"a".repeat(800)}`), |
| fauxAssistantMessage(`recent-history:${"b".repeat(800)}`), |
| fauxAssistantMessage(fauxToolCall("large_result", {}), { stopReason: "toolUse" }), |
| (context) => { |
| order.push("provider"); |
| resumedRequest = JSON.stringify(context.messages); |
| return fauxAssistantMessage("finished after compaction"); |
| }, |
| ]); |
|
|
| await harness.session.prompt("seed old history"); |
| await harness.session.prompt("seed recent history"); |
| const agentStartsBefore = harness.eventsOfType("agent_start").length; |
| await harness.session.prompt("run the large tool"); |
|
|
| expect(order).toEqual(["compaction", "provider"]); |
| expect(observedSettings).toEqual([{ enabled: true, reserveTokens: 400, keepRecentTokens: 1750 }]); |
| expect(harness.eventsOfType("agent_start")).toHaveLength(agentStartsBefore + 1); |
| expect(harness.eventsOfType("compaction_start").at(-1)).toEqual({ |
| type: "compaction_start", |
| reason: "threshold", |
| }); |
| expect(resumedRequest).toContain("compacted history"); |
| expect(resumedRequest).toContain("large-tool-result"); |
| expect(harness.session.getLastAssistantText()).toBe("finished after compaction"); |
| }, |
| ); |
|
|
| it("includes steering queued during compaction in the resumed assistant request", async () => { |
| const largeTool: AgentTool = { |
| name: "large_result", |
| label: "Large result", |
| description: "Returns enough content to cross the compaction threshold", |
| parameters: Type.Object({}), |
| execute: async () => ({ |
| content: [{ type: "text", text: `large-tool-result:${"x".repeat(6800)}` }], |
| details: {}, |
| }), |
| }; |
| let markCompactionStarted = () => {}; |
| const compactionStarted = new Promise<void>((resolve) => { |
| markCompactionStarted = resolve; |
| }); |
| let releaseCompaction = () => {}; |
| const compactionReleased = new Promise<void>((resolve) => { |
| releaseCompaction = resolve; |
| }); |
| const harness = await createHarness({ |
| models: [{ id: "faux-1", contextWindow: 2600, maxTokens: 100 }], |
| settings: { compaction: { enabled: true, reserveTokens: 400, keepRecentTokens: 1750 } }, |
| tools: [largeTool], |
| extensionFactories: [ |
| (pi) => { |
| pi.on("session_before_compact", async (event) => { |
| markCompactionStarted(); |
| await compactionReleased; |
| return { |
| compaction: { |
| summary: "compacted history", |
| firstKeptEntryId: event.preparation.firstKeptEntryId, |
| tokensBefore: event.preparation.tokensBefore, |
| details: {}, |
| }, |
| }; |
| }); |
| }, |
| ], |
| }); |
| harnesses.push(harness); |
| let resumedRequest = ""; |
| harness.setResponses([ |
| fauxAssistantMessage(`old-history:${"a".repeat(800)}`), |
| fauxAssistantMessage(`recent-history:${"b".repeat(800)}`), |
| fauxAssistantMessage(fauxToolCall("large_result", {}), { stopReason: "toolUse" }), |
| (context) => { |
| resumedRequest = JSON.stringify(context.messages); |
| return fauxAssistantMessage("finished after compaction"); |
| }, |
| fauxAssistantMessage("finished after delayed steering"), |
| ]); |
|
|
| await harness.session.prompt("seed old history"); |
| await harness.session.prompt("seed recent history"); |
| const promptPromise = harness.session.prompt("run the large tool"); |
| await compactionStarted; |
| await harness.session.steer("change direction"); |
| releaseCompaction(); |
| await promptPromise; |
|
|
| expect(resumedRequest).toContain("change direction"); |
| expect(harness.faux.state.callCount).toBe(4); |
| }); |
|
|
| it("does not compact after a terminating tool result", async () => { |
| const terminatingTool: AgentTool = { |
| name: "terminate_with_large_result", |
| label: "Terminate with large result", |
| description: "Returns enough content to cross the compaction threshold, then terminates", |
| parameters: Type.Object({}), |
| execute: async () => ({ |
| content: [{ type: "text", text: `large-tool-result:${"x".repeat(6800)}` }], |
| details: {}, |
| terminate: true, |
| }), |
| }; |
| const harness = await createHarness({ |
| models: [{ id: "faux-1", contextWindow: 2600, maxTokens: 100 }], |
| settings: { compaction: { enabled: true, reserveTokens: 400, keepRecentTokens: 1750 } }, |
| tools: [terminatingTool], |
| extensionFactories: [ |
| (pi) => { |
| pi.on("session_before_compact", (event) => ({ |
| compaction: { |
| summary: "unexpected compaction", |
| firstKeptEntryId: event.preparation.firstKeptEntryId, |
| tokensBefore: event.preparation.tokensBefore, |
| details: {}, |
| }, |
| })); |
| }, |
| ], |
| }); |
| harnesses.push(harness); |
| harness.setResponses([ |
| fauxAssistantMessage(`old-history:${"a".repeat(800)}`), |
| fauxAssistantMessage(`recent-history:${"b".repeat(800)}`), |
| fauxAssistantMessage(fauxToolCall("terminate_with_large_result", {}), { stopReason: "toolUse" }), |
| ]); |
|
|
| await harness.session.prompt("seed old history"); |
| await harness.session.prompt("seed recent history"); |
| await harness.session.prompt("run the terminating tool"); |
|
|
| expect(harness.eventsOfType("compaction_start")).toEqual([]); |
| expect(harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction")).toEqual([]); |
| expect(harness.getPendingResponseCount()).toBe(0); |
| }); |
|
|
| it("does not compact when a length stop reaches the desired output limit", async () => { |
| const harness = await createHarness({ |
| models: [{ id: "faux-1", contextWindow: 1_000_000, maxTokens: 100 }], |
| }); |
| harnesses.push(harness); |
| harness.setResponses([fauxAssistantMessage("x".repeat(400), { stopReason: "length" })]); |
|
|
| await harness.session.prompt("hello"); |
|
|
| expect(harness.faux.state.callCount).toBe(1); |
| expect(harness.eventsOfType("compaction_start")).toHaveLength(0); |
| }); |
|
|
| it("stops after one compact-and-retry when a second response is also truncated", async () => { |
| const harness = await createHarness({ |
| models: [{ id: "faux-1", contextWindow: 1_000_000, maxTokens: 100 }], |
| settings: { compaction: { keepRecentTokens: 1, reserveTokens: 0 } }, |
| extensionFactories: [ |
| (pi) => { |
| pi.on("session_before_compact", async (event) => ({ |
| compaction: { |
| summary: "overflow compacted", |
| firstKeptEntryId: event.preparation.firstKeptEntryId, |
| tokensBefore: event.preparation.tokensBefore, |
| details: {}, |
| }, |
| })); |
| }, |
| ], |
| }); |
| harnesses.push(harness); |
| harness.setResponses([ |
| () => fauxAssistantMessage("x".repeat(64), { stopReason: "length", timestamp: Date.now() + 10_000 }), |
| () => fauxAssistantMessage("y".repeat(64), { stopReason: "length", timestamp: Date.now() + 10_000 }), |
| ]); |
|
|
| await harness.session.prompt("x".repeat(5000)); |
|
|
| expect(harness.faux.state.callCount).toBe(2); |
| expect(harness.eventsOfType("compaction_start").filter((event) => event.reason === "overflow")).toHaveLength(1); |
| expect(harness.eventsOfType("compaction_end").at(-1)?.errorMessage).toBe( |
| "Truncated response recovery failed after one compact-and-retry attempt.", |
| ); |
| }); |
|
|
| it("keeps overflow wording when a repeated length stop fills the context window", async () => { |
| const harness = await createHarness({ |
| models: [{ id: "faux-1", contextWindow: 100, maxTokens: 100 }], |
| }); |
| harnesses.push(harness); |
| const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; |
| const lengthOverflowMessage = createAssistant(harness, { |
| stopReason: "length", |
| totalTokens: 100, |
| timestamp: Date.now(), |
| }); |
| const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false); |
| const compactionErrors: string[] = []; |
| harness.session.subscribe((event) => { |
| if (event.type === "compaction_end" && event.errorMessage) { |
| compactionErrors.push(event.errorMessage); |
| } |
| }); |
|
|
| await sessionInternals._checkCompaction(lengthOverflowMessage); |
| await sessionInternals._checkCompaction({ ...lengthOverflowMessage, timestamp: Date.now() + 1 }); |
|
|
| expect(runAutoCompactionSpy).toHaveBeenCalledTimes(1); |
| expect(compactionErrors).toContain( |
| "Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.", |
| ); |
| }); |
|
|
| it("cancels in-progress manual compaction when abortCompaction is called", async () => { |
| const { harness, compactionStarted } = await createAbortableCompactionHarness(); |
| harnesses.push(harness); |
|
|
| const compactPromise = harness.session.compact(); |
| const compactExpectation = expect(compactPromise).rejects.toThrow("Compaction cancelled"); |
| await compactionStarted; |
| harness.session.abortCompaction(); |
|
|
| await compactExpectation; |
| }); |
|
|
| |
| it("aborts an in-progress manual compaction and waits until the session is idle", async () => { |
| const { harness, compactionStarted } = await createAbortableCompactionHarness(); |
| harnesses.push(harness); |
| harness.setResponses([fauxAssistantMessage("continued")]); |
|
|
| const compactPromise = harness.session.compact(); |
| const compactExpectation = expect(compactPromise).rejects.toThrow("Compaction cancelled"); |
| await compactionStarted; |
| await harness.session.abort(); |
| await compactExpectation; |
|
|
| expect(harness.eventsOfType("compaction_end").at(-1)).toMatchObject({ |
| reason: "manual", |
| aborted: true, |
| }); |
| expect(harness.session.isCompacting).toBe(false); |
| expect(harness.session.isIdle).toBe(true); |
|
|
| await expect(harness.session.prompt("next prompt")).resolves.toBeUndefined(); |
| expect(harness.session.getLastAssistantText()).toBe("continued"); |
| }); |
|
|
| it("resumes after threshold compaction when only agent-level queued messages exist", async () => { |
| vi.useFakeTimers(); |
| const harness = await createHarness({ |
| settings: { compaction: { keepRecentTokens: 1 } }, |
| extensionFactories: [ |
| (pi) => { |
| pi.on("session_before_compact", async (event) => ({ |
| compaction: { |
| summary: "auto compacted", |
| firstKeptEntryId: event.preparation.firstKeptEntryId, |
| tokensBefore: event.preparation.tokensBefore, |
| details: {}, |
| }, |
| })); |
| }, |
| ], |
| }); |
| harnesses.push(harness); |
| harness.setResponses([fauxAssistantMessage("one"), fauxAssistantMessage("two")]); |
| await harness.session.prompt("first"); |
| await harness.session.prompt("second"); |
|
|
| harness.session.agent.followUp({ |
| role: "custom", |
| customType: "test", |
| content: [{ type: "text", text: "queued custom" }], |
| display: false, |
| timestamp: Date.now(), |
| }); |
|
|
| const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; |
|
|
| await expect(sessionInternals._runAutoCompaction("threshold", false)).resolves.toBe(true); |
| }); |
|
|
| it("does not retry overflow recovery more than once", async () => { |
| const harness = await createHarness(); |
| harnesses.push(harness); |
| const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; |
| const overflowMessage = createAssistant(harness, { |
| stopReason: "error", |
| errorMessage: "prompt is too long", |
| timestamp: Date.now(), |
| }); |
| const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false); |
| const compactionErrors: string[] = []; |
| harness.session.subscribe((event) => { |
| if (event.type === "compaction_end" && event.errorMessage) { |
| compactionErrors.push(event.errorMessage); |
| } |
| }); |
|
|
| await sessionInternals._checkCompaction(overflowMessage); |
| await sessionInternals._checkCompaction({ ...overflowMessage, timestamp: Date.now() + 1 }); |
|
|
| expect(runAutoCompactionSpy).toHaveBeenCalledTimes(1); |
| expect(compactionErrors).toContain( |
| "Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.", |
| ); |
| }); |
|
|
| it("compacts successful overflow responses without retrying", async () => { |
| const harness = await createHarness({ |
| settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } }, |
| models: [{ id: "faux-1", contextWindow: 1, maxTokens: 100 }], |
| extensionFactories: [ |
| (pi) => { |
| pi.on("session_before_compact", async (event) => ({ |
| compaction: { |
| summary: "successful overflow compacted", |
| firstKeptEntryId: event.preparation.firstKeptEntryId, |
| tokensBefore: event.preparation.tokensBefore, |
| details: {}, |
| }, |
| })); |
| }, |
| ], |
| }); |
| harnesses.push(harness); |
| harness.setResponses([fauxAssistantMessage("completed answer")]); |
|
|
| await expect(harness.session.prompt("hello")).resolves.toBeUndefined(); |
|
|
| const compactionEnd = harness.eventsOfType("compaction_end").at(-1); |
| expect(compactionEnd).toMatchObject({ |
| reason: "overflow", |
| aborted: false, |
| willRetry: false, |
| }); |
| expect(harness.faux.state.callCount).toBe(1); |
| }); |
|
|
| it("ignores stale pre-compaction assistant usage on pre-prompt checks", async () => { |
| const harness = await createHarness(); |
| harnesses.push(harness); |
| const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; |
| const staleTimestamp = Date.now() - 10_000; |
| const staleAssistant = createAssistant(harness, { |
| stopReason: "stop", |
| totalTokens: 610_000, |
| timestamp: staleTimestamp, |
| }); |
|
|
| harness.sessionManager.appendMessage({ |
| role: "user", |
| content: [{ type: "text", text: "before compaction" }], |
| timestamp: staleTimestamp - 1000, |
| }); |
| harness.sessionManager.appendMessage(staleAssistant); |
| const firstKeptEntryId = harness.sessionManager.getEntries()[0]!.id; |
| harness.sessionManager.appendCompaction( |
| "summary", |
| firstKeptEntryId, |
| staleAssistant.usage.totalTokens, |
| undefined, |
| false, |
| ); |
| harness.sessionManager.appendMessage({ |
| role: "user", |
| content: [{ type: "text", text: "after compaction" }], |
| timestamp: Date.now(), |
| }); |
|
|
| const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false); |
|
|
| await sessionInternals._checkCompaction(staleAssistant, false); |
|
|
| expect(runAutoCompactionSpy).not.toHaveBeenCalled(); |
| }); |
|
|
| it("triggers threshold compaction for error messages using the last successful usage", async () => { |
| const harness = await createHarness(); |
| harnesses.push(harness); |
| const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; |
| const successfulAssistant = createAssistant(harness, { |
| stopReason: "stop", |
| totalTokens: 190_000, |
| timestamp: Date.now(), |
| }); |
| const errorAssistant = createAssistant(harness, { |
| stopReason: "error", |
| errorMessage: "529 overloaded", |
| timestamp: Date.now() + 1000, |
| }); |
| harness.session.agent.state.messages = [ |
| { role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now() - 1000 }, |
| successfulAssistant, |
| { role: "user", content: [{ type: "text", text: "retry" }], timestamp: Date.now() + 500 }, |
| errorAssistant, |
| ]; |
|
|
| const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false); |
|
|
| await sessionInternals._checkCompaction(errorAssistant); |
|
|
| expect(runAutoCompactionSpy).toHaveBeenCalledWith("threshold", false); |
| }); |
|
|
| it("does not trigger threshold compaction for error messages when no prior usage exists", async () => { |
| const harness = await createHarness(); |
| harnesses.push(harness); |
| const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; |
| const errorAssistant = createAssistant(harness, { |
| stopReason: "error", |
| errorMessage: "529 overloaded", |
| timestamp: Date.now(), |
| }); |
| harness.session.agent.state.messages = [ |
| { role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now() - 1000 }, |
| errorAssistant, |
| ]; |
|
|
| const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false); |
|
|
| await sessionInternals._checkCompaction(errorAssistant); |
|
|
| expect(runAutoCompactionSpy).not.toHaveBeenCalled(); |
| }); |
|
|
| it("does not trigger threshold compaction when only kept pre-compaction usage exists", async () => { |
| const harness = await createHarness(); |
| harnesses.push(harness); |
| const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; |
| const preCompactionTimestamp = Date.now() - 10_000; |
| const keptAssistant = createAssistant(harness, { |
| stopReason: "stop", |
| totalTokens: 190_000, |
| timestamp: preCompactionTimestamp, |
| }); |
|
|
| harness.sessionManager.appendMessage({ |
| role: "user", |
| content: [{ type: "text", text: "before compaction" }], |
| timestamp: preCompactionTimestamp - 1000, |
| }); |
| harness.sessionManager.appendMessage(keptAssistant); |
| const firstKeptEntryId = harness.sessionManager.getEntries()[0]!.id; |
| harness.sessionManager.appendCompaction( |
| "summary", |
| firstKeptEntryId, |
| keptAssistant.usage.totalTokens, |
| undefined, |
| false, |
| ); |
|
|
| const errorAssistant = createAssistant(harness, { |
| stopReason: "error", |
| errorMessage: "529 overloaded", |
| timestamp: Date.now(), |
| }); |
| harness.session.agent.state.messages = [ |
| { role: "user", content: [{ type: "text", text: "kept user" }], timestamp: preCompactionTimestamp - 1000 }, |
| keptAssistant, |
| { role: "user", content: [{ type: "text", text: "new prompt" }], timestamp: Date.now() - 500 }, |
| errorAssistant, |
| ]; |
|
|
| const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false); |
|
|
| await sessionInternals._checkCompaction(errorAssistant); |
|
|
| expect(runAutoCompactionSpy).not.toHaveBeenCalled(); |
| }); |
|
|
| it("does not trigger threshold compaction below the threshold or when disabled", async () => { |
| const belowThresholdHarness = await createHarness({ |
| settings: { compaction: { enabled: true, reserveTokens: 1000 } }, |
| models: [{ id: "faux-1", contextWindow: 200_000 }], |
| }); |
| harnesses.push(belowThresholdHarness); |
| const disabledHarness = await createHarness({ settings: { compaction: { enabled: false } } }); |
| harnesses.push(disabledHarness); |
|
|
| const belowThresholdInternals = belowThresholdHarness.session as unknown as SessionWithCompactionInternals; |
| const disabledInternals = disabledHarness.session as unknown as SessionWithCompactionInternals; |
| const belowThresholdSpy = vi.spyOn(belowThresholdInternals, "_runAutoCompaction").mockResolvedValue(false); |
| const disabledSpy = vi.spyOn(disabledInternals, "_runAutoCompaction").mockResolvedValue(false); |
|
|
| await belowThresholdInternals._checkCompaction( |
| createAssistant(belowThresholdHarness, { stopReason: "stop", totalTokens: 1_000, timestamp: Date.now() }), |
| ); |
| await disabledInternals._checkCompaction( |
| createAssistant(disabledHarness, { stopReason: "stop", totalTokens: 1_000_000, timestamp: Date.now() }), |
| ); |
|
|
| expect(belowThresholdSpy).not.toHaveBeenCalled(); |
| expect(disabledSpy).not.toHaveBeenCalled(); |
| }); |
| }); |
|
|
