mirror of
https://github.com/earendil-works/pi.git
synced 2026-10-02 08:44:38 +08:00
Armin Ronacher
540e174c72
This adds experimental support for virtual models. A virtual model is a catalog entry that an extension registers with pi.registerVirtualModel(). It does not talk to a provider itself; for every request it picks a physical model and thinking level. The motivation is to make routing policies pluggable: plan on a strong model and implement on a cheap one, pick a model with a classifier, or move to a larger window when the context grows. Today this needs model switching by hand or hacks in prepareRequest. The selection stays virtual: model_change records it, resume restores it, and each response records the physical model and thinking level that produced it. Router state can be stored on the session branch, so it follows /tree and forks and survives compaction. Context limits, compaction and image sizing use the physical model a request was routed to. examples/extensions/jev-router.ts shows a full router that uses the Jev classifier to pick Sol or Terra for planning and hands off to Luna after the first edit. docs/virtual-models.md describes the API.
1019 lines
36 KiB
TypeScript
1019 lines
36 KiB
TypeScript
import type { AgentMessage, AgentTool } from "@earendil-works/pi-agent-core";
|
|
import {
|
|
type AssistantMessage,
|
|
createAssistantMessageEventStream,
|
|
fauxAssistantMessage,
|
|
fauxToolCall,
|
|
getCurrentSystemPrompt,
|
|
getCurrentTools,
|
|
type Model,
|
|
type SimpleStreamOptions,
|
|
type TranscriptContext,
|
|
} 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: TranscriptContext, 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("system");
|
|
expect(harness.session.messages[1]?.role).toBe("compactionSummary");
|
|
});
|
|
|
|
it("checkpoints the replayed system state and folds summarized and retained system patches into it", async () => {
|
|
const harness = await createHarness();
|
|
harnesses.push(harness);
|
|
harness.setResponses([fauxAssistantMessage("declared")]);
|
|
await harness.session.prompt("declare the prompt");
|
|
const declared = harness.session.messages[0];
|
|
if (declared?.role !== "system") throw new Error("expected declared system message");
|
|
|
|
harness.sessionManager.appendMessage({
|
|
role: "system",
|
|
content: "summarized instruction",
|
|
sections: { early: "<early>1</early>" },
|
|
toolsRemoved: [{ name: "bash" }],
|
|
timestamp: Date.now(),
|
|
});
|
|
const firstKeptEntryId = harness.sessionManager.appendMessage({
|
|
role: "user",
|
|
content: [{ type: "text", text: "kept before patch" }],
|
|
timestamp: Date.now(),
|
|
});
|
|
harness.sessionManager.appendMessage({
|
|
role: "system",
|
|
content: "retained instruction",
|
|
sections: { extra: "<extra>late</extra>" },
|
|
toolsRemoved: [{ name: "read" }],
|
|
timestamp: Date.now(),
|
|
});
|
|
harness.sessionManager.appendMessage({
|
|
role: "user",
|
|
content: [{ type: "text", text: "kept after patch" }],
|
|
timestamp: Date.now(),
|
|
});
|
|
harness.sessionManager.appendCompaction("compacted", firstKeptEntryId, 100);
|
|
|
|
const messages = harness.sessionManager.buildSessionContext().messages;
|
|
expect(messages.map((message) => message.role)).toEqual(["system", "compactionSummary", "user", "user"]);
|
|
const checkpoint = messages[0];
|
|
if (checkpoint?.role !== "system") throw new Error("expected checkpoint system message");
|
|
expect(checkpoint.content).toBe("summarized instruction\n\nretained instruction");
|
|
expect(checkpoint.sections).toEqual({
|
|
...declared.sections,
|
|
early: "<early>1</early>",
|
|
extra: "<extra>late</extra>",
|
|
});
|
|
expect(checkpoint.toolsAdded?.map((tool) => tool.name)).toEqual(
|
|
harness.session.getActiveToolNames().filter((name) => name !== "read" && name !== "bash"),
|
|
);
|
|
});
|
|
|
|
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);
|
|
// Auth is resolved only when Pi summarizes itself, after checking there is something to compact.
|
|
seedCompactableSession(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);
|
|
const summaryResponse = (_context: TranscriptContext, options: SimpleStreamOptions | undefined) => {
|
|
expect(options?.apiKey).toBeUndefined();
|
|
expect(options?.headers).toEqual({ Authorization: "Bearer ambient-token" });
|
|
return fauxAssistantMessage("summary with bearer auth");
|
|
};
|
|
harness.setResponses([summaryResponse]);
|
|
|
|
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: TranscriptContext | undefined;
|
|
let requestOptions: SimpleStreamOptions | undefined;
|
|
useSummaryStreamFn(harness, "standalone summary", (context, options) => {
|
|
requestContext = context;
|
|
requestOptions = options;
|
|
});
|
|
|
|
await harness.session.compact();
|
|
|
|
expect(transformContext).not.toHaveBeenCalled();
|
|
expect(getCurrentSystemPrompt(requestContext?.messages ?? [])).not.toBe(harness.session.agent.state.systemPrompt);
|
|
expect(getCurrentTools(requestContext?.messages ?? [])).toEqual([]);
|
|
// Regression test for #9652: split-turn summaries use a clear Markdown conversation boundary.
|
|
expect(JSON.stringify(requestContext?.messages)).toContain("# Conversation\\n[User]: message to compact");
|
|
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");
|
|
});
|
|
|
|
// Regression coverage for #8133: model overrides must also apply between assistant turns.
|
|
// Regression coverage for #9740: an oversized trailing tool result must still produce a cut point.
|
|
it.each([false, true])(
|
|
"compacts after an oversized tool result in the same run (model override: %s)",
|
|
async (modelOverride) => {
|
|
const toolResult = `large-tool-result:${"x".repeat(8000)}`;
|
|
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.slice(0, 2)).toEqual(["compaction", "provider"]);
|
|
expect(observedSettings[0]).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;
|
|
});
|
|
|
|
// Regression test for #8920.
|
|
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();
|
|
});
|
|
});
|