mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
feat(coding-agent): add cache-friendly compaction primitives
Add provider-context prefix preparation and optional summary request settings needed by cache-friendly compaction. Existing AgentSession callers continue to pass undefined, so standalone summarization remains the default until the feature is wired in.
This commit is contained in:
Vegard Stikbakke
parent
ef8dc7385b
commit
cff1cf52c6
7 changed files with 528 additions and 45 deletions
|
|
@ -37,6 +37,7 @@ export * from "./providers/faux.ts";
|
|||
export * from "./session-resources.ts";
|
||||
export * from "./types.ts";
|
||||
export * from "./utils/diagnostics.ts";
|
||||
export { estimateContextTokens } from "./utils/estimate.ts";
|
||||
export * from "./utils/event-stream.ts";
|
||||
export * from "./utils/json-parse.ts";
|
||||
export * from "./utils/overflow.ts";
|
||||
|
|
|
|||
|
|
@ -281,7 +281,9 @@ pi.on("session_before_compact", async (event, ctx) => {
|
|||
const { preparation, branchEntries, customInstructions, reason, willRetry, signal } = event;
|
||||
|
||||
// preparation.messagesToSummarize - messages to summarize
|
||||
// preparation.sourceMessages - optional active-context history prefix, including any previous summary
|
||||
// preparation.turnPrefixMessages - split turn prefix (if isSplitTurn)
|
||||
// preparation.turnPrefixSourceMessages - optional active-context prefix through the split turn
|
||||
// preparation.previousSummary - previous compaction summary
|
||||
// preparation.fileOps - extracted file operations
|
||||
// preparation.tokensBefore - context tokens before compaction
|
||||
|
|
|
|||
|
|
@ -1813,7 +1813,7 @@ export class AgentSession {
|
|||
env,
|
||||
this.settingsManager.getRetrySettings(),
|
||||
this._summarizationRetryCallbacks({ source: "compaction", reason }),
|
||||
undefined, // sessionId
|
||||
undefined, // cacheFriendly
|
||||
);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,14 @@
|
|||
*/
|
||||
|
||||
import type { AgentMessage, StreamFn, ThinkingLevel } from "@earendil-works/pi-agent-core";
|
||||
import { contentText, type RetryCallbacks, type RetryPolicy, retryAssistantCall, uuidv7 } from "@earendil-works/pi-ai";
|
||||
import {
|
||||
contentText,
|
||||
estimateContextTokens as estimateProviderContextTokens,
|
||||
type RetryCallbacks,
|
||||
type RetryPolicy,
|
||||
retryAssistantCall,
|
||||
uuidv7,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import type { AssistantMessage, Context, Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat";
|
||||
import { completeSimple } from "@earendil-works/pi-ai/compat";
|
||||
import { convertToLlm } from "../messages.ts";
|
||||
|
|
@ -95,6 +102,27 @@ function getMessageFromEntryForCompaction(entry: SessionEntry): AgentMessage | u
|
|||
return sessionEntryToContextMessages(entry)[0];
|
||||
}
|
||||
|
||||
/** Build an active-context prefix, placing the previous compaction summary before its retained messages. */
|
||||
function collectSourceMessages(
|
||||
entries: SessionEntry[],
|
||||
startIndex: number,
|
||||
endIndex: number,
|
||||
previousCompactionIndex: number,
|
||||
): AgentMessage[] {
|
||||
const messages: AgentMessage[] = [];
|
||||
if (previousCompactionIndex >= 0) {
|
||||
messages.push(...sessionEntryToContextMessages(entries[previousCompactionIndex]));
|
||||
}
|
||||
for (let i = startIndex; i < endIndex; i++) {
|
||||
// The latest compaction was moved to the front above. Keep older compaction
|
||||
// entries because buildSessionContext retains them in the active provider prefix.
|
||||
if (i !== previousCompactionIndex) {
|
||||
messages.push(...sessionEntryToContextMessages(entries[i]));
|
||||
}
|
||||
}
|
||||
return messages;
|
||||
}
|
||||
|
||||
/** Result from compact() - SessionManager adds uuid/parentUuid when saving */
|
||||
export interface CompactionResult<T = unknown> {
|
||||
summary: string;
|
||||
|
|
@ -140,6 +168,19 @@ export interface CompactionSettings {
|
|||
keepRecentTokens: number;
|
||||
}
|
||||
|
||||
/** Active provider contexts and request settings used to preserve cacheable compaction prefixes. */
|
||||
export interface CacheFriendlySummaryOptions {
|
||||
/** Exact provider context prefix containing the history to summarize. */
|
||||
sourceContext?: Context;
|
||||
/** Exact provider context prefix containing a split turn's prefix. */
|
||||
turnPrefixSourceContext?: Context;
|
||||
/** Provider request settings copied from the active agent request path. */
|
||||
requestOptions?: Pick<
|
||||
SimpleStreamOptions,
|
||||
"sessionId" | "onPayload" | "onResponse" | "transport" | "thinkingBudgets" | "maxRetryDelayMs"
|
||||
>;
|
||||
}
|
||||
|
||||
export const DEFAULT_COMPACTION_SETTINGS: CompactionSettings = {
|
||||
enabled: true,
|
||||
reserveTokens: 16384,
|
||||
|
|
@ -549,6 +590,10 @@ const UPDATE_SUMMARIZATION_PROMPT = `The messages above are NEW conversation mes
|
|||
|
||||
${UPDATE_SUMMARIZATION_INSTRUCTIONS}`;
|
||||
|
||||
const SOURCE_CONTEXT_UPDATE_SUMMARIZATION_PROMPT = `The messages above contain an existing structured summary of earlier conversation history followed by NEW conversation messages.
|
||||
|
||||
${UPDATE_SUMMARIZATION_INSTRUCTIONS}`;
|
||||
|
||||
function createSummarizationOptions(
|
||||
model: Model<any>,
|
||||
maxTokens: number,
|
||||
|
|
@ -557,9 +602,18 @@ function createSummarizationOptions(
|
|||
env: Record<string, string> | undefined,
|
||||
signal: AbortSignal | undefined,
|
||||
thinkingLevel: ThinkingLevel | undefined,
|
||||
sessionId: string | undefined,
|
||||
requestOptions: CacheFriendlySummaryOptions["requestOptions"] | undefined,
|
||||
cacheRetention: SimpleStreamOptions["cacheRetention"] | undefined,
|
||||
): SimpleStreamOptions {
|
||||
const options: SimpleStreamOptions = { maxTokens, signal, apiKey, headers, env, sessionId };
|
||||
const options: SimpleStreamOptions = {
|
||||
...requestOptions,
|
||||
maxTokens,
|
||||
signal,
|
||||
apiKey,
|
||||
headers,
|
||||
env,
|
||||
cacheRetention,
|
||||
};
|
||||
const refusalFallbacks = getAnthropicSummarizationFallback(model);
|
||||
if (refusalFallbacks) {
|
||||
options.refusalFallbacks = refusalFallbacks;
|
||||
|
|
@ -585,12 +639,15 @@ export async function completeSummarization(
|
|||
retry?: RetryPolicy,
|
||||
callbacks?: RetryCallbacks,
|
||||
): Promise<AssistantMessage> {
|
||||
// Avoid cache writes for one-off summaries. Reuse caller-supplied routing when available;
|
||||
// callers without a session ID, including branch summaries, receive a fresh routing ID.
|
||||
// Standalone one-off summaries default to no prompt caching. Cache-friendly callers
|
||||
// explicitly request short retention so providers can reuse the active prefix.
|
||||
// Callers without a session ID, including standalone branch summaries, receive a fresh routing ID.
|
||||
const requestOptions: SimpleStreamOptions = {
|
||||
...options,
|
||||
cacheRetention: "none",
|
||||
cacheRetention: options.cacheRetention ?? "none",
|
||||
sessionId: options.sessionId ?? uuidv7(),
|
||||
// Anthropic invalidates the messages cache when tool_choice changes. Its 20-content-block lookup
|
||||
// can still reuse an earlier user-message checkpoint, but long tool-heavy turns may need reprocessing.
|
||||
toolChoice: "none",
|
||||
};
|
||||
const produce = async (): Promise<AssistantMessage> =>
|
||||
|
|
@ -618,7 +675,7 @@ export async function generateSummary(
|
|||
env?: Record<string, string>,
|
||||
retry?: RetryPolicy,
|
||||
callbacks?: RetryCallbacks,
|
||||
sessionId?: string,
|
||||
cacheFriendly?: Pick<CacheFriendlySummaryOptions, "sourceContext" | "requestOptions">,
|
||||
): Promise<string> {
|
||||
return (
|
||||
await generateSummaryWithUsage(
|
||||
|
|
@ -635,25 +692,47 @@ export async function generateSummary(
|
|||
env,
|
||||
retry,
|
||||
callbacks,
|
||||
sessionId,
|
||||
cacheFriendly,
|
||||
)
|
||||
).text;
|
||||
}
|
||||
|
||||
/** Build the provider context for a standalone summary request. */
|
||||
function buildSummarizationContext(promptText: string): Context {
|
||||
/** Build a standalone summary request or append its instruction to an existing provider context. */
|
||||
function buildSummarizationContext(promptText: string, sourceContext?: Context): Context {
|
||||
const instructionMessage = {
|
||||
role: "user" as const,
|
||||
content: [{ type: "text" as const, text: promptText }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
if (sourceContext) {
|
||||
return {
|
||||
...sourceContext,
|
||||
messages: [...sourceContext.messages, instructionMessage],
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
systemPrompt: SUMMARIZATION_SYSTEM_PROMPT,
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: [{ type: "text", text: promptText }],
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
messages: [instructionMessage],
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Extra room for provider framing and tokenizer variance omitted by the heuristic context estimate.
|
||||
* This matches the 4096-token margin used when normal simple requests clamp maxTokens to their context window.
|
||||
*/
|
||||
const CACHE_FRIENDLY_CONTEXT_SAFETY_TOKENS = 4096;
|
||||
|
||||
/** Whether the source context leaves room for the requested summary output and provider safety margin. */
|
||||
function cacheFriendlyContextFits(model: Model<any>, context: Context, maxTokens: number): boolean {
|
||||
return (
|
||||
model.contextWindow <= 0 ||
|
||||
estimateProviderContextTokens(context).tokens + maxTokens + CACHE_FRIENDLY_CONTEXT_SAFETY_TOKENS <=
|
||||
model.contextWindow
|
||||
);
|
||||
}
|
||||
|
||||
/** Generate or update a conversation summary and return its provider usage. */
|
||||
export async function generateSummaryWithUsage(
|
||||
currentMessages: AgentMessage[],
|
||||
|
|
@ -669,31 +748,55 @@ export async function generateSummaryWithUsage(
|
|||
env?: Record<string, string>,
|
||||
retry?: RetryPolicy,
|
||||
callbacks?: RetryCallbacks,
|
||||
sessionId?: string,
|
||||
cacheFriendly?: Pick<CacheFriendlySummaryOptions, "sourceContext" | "requestOptions">,
|
||||
): Promise<{ text: string; usage: Usage }> {
|
||||
const maxTokens = Math.min(
|
||||
Math.floor(0.8 * reserveTokens),
|
||||
model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY,
|
||||
);
|
||||
// Provider-visible history prefix to reuse instead of serializing messages into a standalone prompt.
|
||||
let sourceContext = cacheFriendly?.sourceContext;
|
||||
|
||||
// Use update prompt if we have a previous summary, otherwise initial prompt
|
||||
let basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT;
|
||||
// Cache-friendly source contexts already contain the previous compaction summary,
|
||||
// but still need iterative-update instructions so prior information is preserved.
|
||||
let basePrompt = previousSummary
|
||||
? sourceContext
|
||||
? SOURCE_CONTEXT_UPDATE_SUMMARIZATION_PROMPT
|
||||
: UPDATE_SUMMARIZATION_PROMPT
|
||||
: SUMMARIZATION_PROMPT;
|
||||
if (customInstructions) {
|
||||
basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`;
|
||||
}
|
||||
|
||||
// Serialize conversation to text so model doesn't try to continue it
|
||||
// Convert to LLM messages first (handles custom types like bashExecution, custom, etc.)
|
||||
const llmMessages = convertToLlm(currentMessages);
|
||||
const conversationText = serializeConversation(llmMessages);
|
||||
// A valid prefix can still be too large to leave the intended summary budget,
|
||||
// especially during overflow recovery or after switching to a smaller-context model.
|
||||
// In that case, use standalone serialization, which truncates large tool results.
|
||||
if (
|
||||
sourceContext &&
|
||||
!cacheFriendlyContextFits(model, buildSummarizationContext(basePrompt, sourceContext), maxTokens)
|
||||
) {
|
||||
sourceContext = undefined;
|
||||
basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT;
|
||||
if (customInstructions) {
|
||||
basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`;
|
||||
}
|
||||
}
|
||||
|
||||
// Build the prompt with conversation wrapped in tags
|
||||
let promptText = `<conversation>\n${conversationText}\n</conversation>\n\n`;
|
||||
if (previousSummary) {
|
||||
// Source contexts already contain the conversation. Standalone requests serialize it
|
||||
// so the model treats those messages as data rather than continuing the conversation.
|
||||
let promptText = "";
|
||||
if (!sourceContext) {
|
||||
const llmMessages = convertToLlm(currentMessages);
|
||||
const conversationText = serializeConversation(llmMessages);
|
||||
promptText = `<conversation>\n${conversationText}\n</conversation>\n\n`;
|
||||
}
|
||||
if (previousSummary && !sourceContext) {
|
||||
promptText += `<previous-summary>\n${previousSummary}\n</previous-summary>\n\n`;
|
||||
}
|
||||
promptText += basePrompt;
|
||||
|
||||
const sourceRequestOptions = sourceContext ? cacheFriendly?.requestOptions : undefined;
|
||||
const sourceCacheRetention = sourceContext ? "short" : undefined;
|
||||
const completionOptions = createSummarizationOptions(
|
||||
model,
|
||||
maxTokens,
|
||||
|
|
@ -702,12 +805,13 @@ export async function generateSummaryWithUsage(
|
|||
env,
|
||||
signal,
|
||||
thinkingLevel,
|
||||
sessionId,
|
||||
sourceRequestOptions,
|
||||
sourceCacheRetention,
|
||||
);
|
||||
|
||||
const response = await completeSummarization(
|
||||
model,
|
||||
buildSummarizationContext(promptText),
|
||||
buildSummarizationContext(promptText, sourceContext),
|
||||
completionOptions,
|
||||
streamFn,
|
||||
retry,
|
||||
|
|
@ -735,8 +839,15 @@ export interface CompactionPreparation {
|
|||
firstKeptEntryId: string;
|
||||
/** Messages that will be summarized and discarded */
|
||||
messagesToSummarize: AgentMessage[];
|
||||
/**
|
||||
* Active-context prefix for the history summary.
|
||||
* Includes the previous compaction summary before messages retained by that compaction.
|
||||
*/
|
||||
sourceMessages?: AgentMessage[];
|
||||
/** Messages that will be turned into turn prefix summary (if splitting) */
|
||||
turnPrefixMessages: AgentMessage[];
|
||||
/** Active-context prefix through the split-turn prefix, or empty when not splitting. */
|
||||
turnPrefixSourceMessages?: AgentMessage[];
|
||||
/** Whether this is a split turn (cut point in middle of turn) */
|
||||
isSplitTurn: boolean;
|
||||
tokensBefore: number;
|
||||
|
|
@ -794,6 +905,8 @@ export function prepareCompaction(
|
|||
if (msg) messagesToSummarize.push(msg);
|
||||
}
|
||||
|
||||
const sourceMessages = collectSourceMessages(pathEntries, boundaryStart, historyEnd, prevCompactionIndex);
|
||||
|
||||
// Messages for turn prefix summary (if splitting a turn)
|
||||
const turnPrefixMessages: AgentMessage[] = [];
|
||||
if (cutPoint.isSplitTurn) {
|
||||
|
|
@ -802,6 +915,9 @@ export function prepareCompaction(
|
|||
if (msg) turnPrefixMessages.push(msg);
|
||||
}
|
||||
}
|
||||
const turnPrefixSourceMessages = cutPoint.isSplitTurn
|
||||
? collectSourceMessages(pathEntries, boundaryStart, cutPoint.firstKeptEntryIndex, prevCompactionIndex)
|
||||
: [];
|
||||
|
||||
if (messagesToSummarize.length === 0 && turnPrefixMessages.length === 0) {
|
||||
return undefined;
|
||||
|
|
@ -820,7 +936,9 @@ export function prepareCompaction(
|
|||
return {
|
||||
firstKeptEntryId,
|
||||
messagesToSummarize,
|
||||
sourceMessages,
|
||||
turnPrefixMessages,
|
||||
turnPrefixSourceMessages,
|
||||
isSplitTurn: cutPoint.isSplitTurn,
|
||||
tokensBefore,
|
||||
previousSummary,
|
||||
|
|
@ -848,13 +966,30 @@ Summarize the prefix to provide context for the retained suffix:
|
|||
|
||||
Be concise. Focus on what's needed to understand the kept suffix.`;
|
||||
|
||||
const SOURCE_CONTEXT_TURN_PREFIX_SUMMARIZATION_PROMPT = `The final turn in the source conversation was too large to keep in full. Its SUFFIX (recent work) is retained.
|
||||
|
||||
The source conversation may also contain complete earlier turns for background. Summarize only the final, incomplete turn. It begins with the last user-role request before this instruction. Do not summarize earlier turns except for details needed to understand this final turn's prefix.
|
||||
|
||||
Summarize the prefix to provide context for the retained suffix:
|
||||
|
||||
## Original Request
|
||||
[What did the user ask for in this turn?]
|
||||
|
||||
## Early Progress
|
||||
- [Key decisions and work done in the prefix]
|
||||
|
||||
## Context for Suffix
|
||||
- [Information needed to understand the retained recent work]
|
||||
|
||||
Be concise. Focus on what's needed to understand the kept suffix.`;
|
||||
|
||||
/**
|
||||
* Generate summaries for compaction using prepared data.
|
||||
* Returns CompactionResult - SessionManager adds uuid/parentUuid when saving.
|
||||
*
|
||||
* @param preparation - Pre-calculated preparation from prepareCompaction()
|
||||
* @param customInstructions - Optional custom focus for the summary
|
||||
* @param sessionId - Optional routing session ID forwarded without enabling prompt caching
|
||||
* @param cacheFriendly - Active provider contexts and request settings for cache-friendly summarization
|
||||
*/
|
||||
export async function compact(
|
||||
preparation: CompactionPreparation,
|
||||
|
|
@ -868,7 +1003,7 @@ export async function compact(
|
|||
env?: Record<string, string>,
|
||||
retry?: RetryPolicy,
|
||||
callbacks?: RetryCallbacks,
|
||||
sessionId?: string,
|
||||
cacheFriendly?: CacheFriendlySummaryOptions,
|
||||
): Promise<CompactionResult> {
|
||||
const {
|
||||
firstKeptEntryId,
|
||||
|
|
@ -903,7 +1038,10 @@ export async function compact(
|
|||
env,
|
||||
retry,
|
||||
callbacks,
|
||||
sessionId,
|
||||
{
|
||||
sourceContext: cacheFriendly?.sourceContext,
|
||||
requestOptions: cacheFriendly?.requestOptions,
|
||||
},
|
||||
);
|
||||
historyText = historyResult.text;
|
||||
historyUsage = historyResult.usage;
|
||||
|
|
@ -920,7 +1058,10 @@ export async function compact(
|
|||
streamFn,
|
||||
retry,
|
||||
callbacks,
|
||||
sessionId,
|
||||
{
|
||||
sourceContext: cacheFriendly?.turnPrefixSourceContext,
|
||||
requestOptions: cacheFriendly?.requestOptions,
|
||||
},
|
||||
);
|
||||
// Merge into single summary
|
||||
summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.text}`;
|
||||
|
|
@ -941,7 +1082,10 @@ export async function compact(
|
|||
env,
|
||||
retry,
|
||||
callbacks,
|
||||
sessionId,
|
||||
{
|
||||
sourceContext: cacheFriendly?.sourceContext,
|
||||
requestOptions: cacheFriendly?.requestOptions,
|
||||
},
|
||||
);
|
||||
summary = result.text;
|
||||
summaryUsage = result.usage;
|
||||
|
|
@ -979,20 +1123,51 @@ async function generateTurnPrefixSummary(
|
|||
streamFn?: StreamFn,
|
||||
retry?: RetryPolicy,
|
||||
callbacks?: RetryCallbacks,
|
||||
sessionId?: string,
|
||||
cacheFriendly?: Pick<CacheFriendlySummaryOptions, "sourceContext" | "requestOptions">,
|
||||
): Promise<{ text: string; usage: Usage }> {
|
||||
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}`;
|
||||
|
||||
// Reuse the provider-visible split-turn prefix only when it leaves room for the summary;
|
||||
// otherwise serialize and truncate the messages in a standalone prompt.
|
||||
let sourceContext = cacheFriendly?.sourceContext;
|
||||
if (
|
||||
sourceContext &&
|
||||
!cacheFriendlyContextFits(
|
||||
model,
|
||||
buildSummarizationContext(SOURCE_CONTEXT_TURN_PREFIX_SUMMARIZATION_PROMPT, sourceContext),
|
||||
maxTokens,
|
||||
)
|
||||
) {
|
||||
sourceContext = undefined;
|
||||
}
|
||||
let promptText: string;
|
||||
if (sourceContext) {
|
||||
promptText = SOURCE_CONTEXT_TURN_PREFIX_SUMMARIZATION_PROMPT;
|
||||
} else {
|
||||
const llmMessages = convertToLlm(messages);
|
||||
const conversationText = serializeConversation(llmMessages);
|
||||
promptText = `<conversation>\n${conversationText}\n</conversation>\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`;
|
||||
}
|
||||
|
||||
const sourceRequestOptions = sourceContext ? cacheFriendly?.requestOptions : undefined;
|
||||
const sourceCacheRetention = sourceContext ? "short" : undefined;
|
||||
const response = await completeSummarization(
|
||||
model,
|
||||
buildSummarizationContext(promptText),
|
||||
createSummarizationOptions(model, maxTokens, apiKey, headers, env, signal, thinkingLevel, sessionId),
|
||||
buildSummarizationContext(promptText, sourceContext),
|
||||
createSummarizationOptions(
|
||||
model,
|
||||
maxTokens,
|
||||
apiKey,
|
||||
headers,
|
||||
env,
|
||||
signal,
|
||||
thinkingLevel,
|
||||
sourceRequestOptions,
|
||||
sourceCacheRetention,
|
||||
),
|
||||
streamFn,
|
||||
retry,
|
||||
callbacks,
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ export { readStoredCredential } from "./core/auth-storage.ts";
|
|||
export {
|
||||
type BranchPreparation,
|
||||
type BranchSummaryResult,
|
||||
type CacheFriendlySummaryOptions,
|
||||
type CollectEntriesResult,
|
||||
type CompactionResult,
|
||||
type CutPointResult,
|
||||
|
|
|
|||
|
|
@ -115,7 +115,7 @@ describe("generateSummary reasoning options", () => {
|
|||
expect(sessionIds[0]).not.toBe(sessionIds[1]);
|
||||
});
|
||||
|
||||
it("honors a caller-supplied routing session without prompt caching", async () => {
|
||||
it("honors caller-supplied cache retention and routing", async () => {
|
||||
await completeSummarization(
|
||||
createModel(false),
|
||||
{ systemPrompt: "Summarize", messages: [] },
|
||||
|
|
@ -124,7 +124,7 @@ describe("generateSummary reasoning options", () => {
|
|||
|
||||
expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({
|
||||
sessionId: "current-routing-session",
|
||||
cacheRetention: "none",
|
||||
cacheRetention: "long",
|
||||
toolChoice: "none",
|
||||
});
|
||||
});
|
||||
|
|
@ -148,6 +148,229 @@ describe("generateSummary reasoning options", () => {
|
|||
expect(prompt).toContain("<conversation>");
|
||||
});
|
||||
|
||||
it("appends instructions to a cache-friendly source context", async () => {
|
||||
const sourceContext: Context = {
|
||||
systemPrompt: "You are a coding agent.",
|
||||
messages: [{ role: "user", content: "Previous summary and original request", timestamp: 1 }],
|
||||
tools: [],
|
||||
};
|
||||
const onPayload = async (payload: unknown) => payload;
|
||||
const onResponse = async () => {};
|
||||
const preparation: CompactionPreparation = {
|
||||
firstKeptEntryId: "entry-keep",
|
||||
messagesToSummarize: messages,
|
||||
turnPrefixMessages: [],
|
||||
isSplitTurn: false,
|
||||
tokensBefore: 100,
|
||||
previousSummary: "Previous summary",
|
||||
fileOps: { read: new Set(), written: new Set(), edited: new Set() },
|
||||
settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 },
|
||||
};
|
||||
|
||||
await compact(
|
||||
preparation,
|
||||
createModel(false),
|
||||
"test-key",
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
{
|
||||
sourceContext,
|
||||
requestOptions: {
|
||||
sessionId: "routing-session",
|
||||
onPayload,
|
||||
onResponse,
|
||||
transport: "websocket",
|
||||
thinkingBudgets: { low: 1234 },
|
||||
maxRetryDelayMs: 4321,
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
const requestContext = completeSimpleMock.mock.calls[0][1] as Context;
|
||||
expect(requestContext.systemPrompt).toBe(sourceContext.systemPrompt);
|
||||
expect(requestContext.tools).toBe(sourceContext.tools);
|
||||
expect(requestContext.messages.slice(0, -1)).toEqual(sourceContext.messages);
|
||||
const instruction = JSON.stringify(requestContext.messages.at(-1));
|
||||
expect(instruction).toContain("existing structured summary of earlier conversation history");
|
||||
expect(instruction).toContain("PRESERVE all existing information from the previous summary");
|
||||
expect(instruction).not.toContain("<conversation>");
|
||||
expect(instruction).not.toContain("<previous-summary>");
|
||||
expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({
|
||||
cacheRetention: "short",
|
||||
sessionId: "routing-session",
|
||||
toolChoice: "none",
|
||||
transport: "websocket",
|
||||
thinkingBudgets: { low: 1234 },
|
||||
maxRetryDelayMs: 4321,
|
||||
});
|
||||
expect(completeSimpleMock.mock.calls[0][2]?.onPayload).toBe(onPayload);
|
||||
expect(completeSimpleMock.mock.calls[0][2]?.onResponse).toBe(onResponse);
|
||||
expect(sourceContext.messages).toHaveLength(1);
|
||||
});
|
||||
|
||||
it("falls back to standalone summarization when a cache-friendly source cannot leave the summary budget", async () => {
|
||||
const oversizedToolResult: AgentMessage = {
|
||||
role: "toolResult",
|
||||
toolCallId: "tool-call-1",
|
||||
toolName: "read",
|
||||
content: [{ type: "text", text: "x".repeat(800_000) }],
|
||||
isError: false,
|
||||
timestamp: 1,
|
||||
};
|
||||
const sourceContext: Context = {
|
||||
systemPrompt: "You are a coding agent.",
|
||||
messages: [oversizedToolResult],
|
||||
tools: [],
|
||||
};
|
||||
const preparation: CompactionPreparation = {
|
||||
firstKeptEntryId: "entry-keep",
|
||||
messagesToSummarize: [oversizedToolResult],
|
||||
turnPrefixMessages: [],
|
||||
isSplitTurn: false,
|
||||
tokensBefore: 250_000,
|
||||
fileOps: { read: new Set(), written: new Set(), edited: new Set() },
|
||||
settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 },
|
||||
};
|
||||
|
||||
await compact(
|
||||
preparation,
|
||||
createModel(false),
|
||||
"test-key",
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
{ sourceContext, requestOptions: { sessionId: "routing-session" } },
|
||||
);
|
||||
|
||||
const requestContext = completeSimpleMock.mock.calls[0][1] as Context;
|
||||
const prompt = JSON.stringify(requestContext.messages);
|
||||
expect(requestContext.systemPrompt).not.toBe(sourceContext.systemPrompt);
|
||||
expect(requestContext.tools).toBeUndefined();
|
||||
expect(prompt).toContain("<conversation>");
|
||||
expect(prompt).toContain("more characters truncated");
|
||||
expect(prompt).not.toContain("x".repeat(3000));
|
||||
expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({ cacheRetention: "none", toolChoice: "none" });
|
||||
expect(completeSimpleMock.mock.calls[0][2]?.sessionId).not.toBe("routing-session");
|
||||
});
|
||||
|
||||
it("falls back to standalone summarization for an oversized split-turn source", async () => {
|
||||
const splitUser = { role: "user" as const, content: "Large final request", timestamp: 1 };
|
||||
const oversizedToolResult: AgentMessage = {
|
||||
role: "toolResult",
|
||||
toolCallId: "tool-call-1",
|
||||
toolName: "read",
|
||||
content: [{ type: "text", text: "x".repeat(800_000) }],
|
||||
isError: false,
|
||||
timestamp: 2,
|
||||
};
|
||||
const turnPrefixSourceContext: Context = {
|
||||
systemPrompt: "You are a coding agent.",
|
||||
messages: [splitUser, oversizedToolResult],
|
||||
tools: [],
|
||||
};
|
||||
const preparation: CompactionPreparation = {
|
||||
firstKeptEntryId: "entry-keep",
|
||||
messagesToSummarize: [],
|
||||
turnPrefixMessages: [splitUser, oversizedToolResult],
|
||||
isSplitTurn: true,
|
||||
tokensBefore: 250_000,
|
||||
fileOps: { read: new Set(), written: new Set(), edited: new Set() },
|
||||
settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 },
|
||||
};
|
||||
|
||||
await compact(
|
||||
preparation,
|
||||
createModel(false),
|
||||
"test-key",
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
{ turnPrefixSourceContext, requestOptions: { sessionId: "routing-session" } },
|
||||
);
|
||||
|
||||
const requestContext = completeSimpleMock.mock.calls[0][1] as Context;
|
||||
const prompt = JSON.stringify(requestContext.messages);
|
||||
expect(requestContext.systemPrompt).not.toBe(turnPrefixSourceContext.systemPrompt);
|
||||
expect(requestContext.tools).toBeUndefined();
|
||||
expect(prompt).toContain("This is the PREFIX of a turn that was too large to keep");
|
||||
expect(prompt).toContain("more characters truncated");
|
||||
expect(prompt).not.toContain("source conversation may also contain complete earlier turns");
|
||||
expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({ cacheRetention: "none", toolChoice: "none" });
|
||||
expect(completeSimpleMock.mock.calls[0][2]?.sessionId).not.toBe("routing-session");
|
||||
});
|
||||
|
||||
it("limits a cache-friendly split-turn summary to the final incomplete turn", async () => {
|
||||
const earlierUser = { role: "user" as const, content: "Earlier request", timestamp: 1 };
|
||||
const earlierAssistant: AssistantMessage = {
|
||||
...mockSummaryResponse,
|
||||
content: [{ type: "text", text: "Earlier work completed" }],
|
||||
timestamp: 2,
|
||||
};
|
||||
const splitUser = { role: "user" as const, content: "Large final request", timestamp: 3 };
|
||||
const earlyAssistant: AssistantMessage = {
|
||||
...mockSummaryResponse,
|
||||
content: [{ type: "text", text: "Early work in final turn" }],
|
||||
timestamp: 4,
|
||||
};
|
||||
const turnPrefixSourceContext: Context = {
|
||||
systemPrompt: "You are a coding agent.",
|
||||
messages: [earlierUser, earlierAssistant, splitUser, earlyAssistant],
|
||||
tools: [],
|
||||
};
|
||||
const preparation: CompactionPreparation = {
|
||||
firstKeptEntryId: "entry-keep",
|
||||
messagesToSummarize: [],
|
||||
turnPrefixMessages: [splitUser, earlyAssistant],
|
||||
isSplitTurn: true,
|
||||
tokensBefore: 100,
|
||||
fileOps: { read: new Set(), written: new Set(), edited: new Set() },
|
||||
settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 },
|
||||
};
|
||||
|
||||
await compact(
|
||||
preparation,
|
||||
createModel(false),
|
||||
"test-key",
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
{ turnPrefixSourceContext },
|
||||
);
|
||||
|
||||
const requestContext = completeSimpleMock.mock.calls[0][1] as Context;
|
||||
expect(requestContext.messages.slice(0, -1)).toEqual(turnPrefixSourceContext.messages);
|
||||
const instruction = JSON.stringify(requestContext.messages.at(-1));
|
||||
expect(instruction).toContain("Summarize only the final, incomplete turn");
|
||||
expect(instruction).toContain("last user-role request before this instruction");
|
||||
expect(instruction).toContain("Do not summarize earlier turns");
|
||||
expect(instruction).not.toContain("Earlier request");
|
||||
expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({
|
||||
cacheRetention: "short",
|
||||
toolChoice: "none",
|
||||
});
|
||||
});
|
||||
|
||||
it("rejects tool calls from conversation summaries", async () => {
|
||||
completeSimpleMock.mockResolvedValueOnce(mockToolCallResponse);
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import { readFileSync } from "fs";
|
|||
import { join } from "path";
|
||||
import { beforeEach, describe, expect, it } from "vitest";
|
||||
import {
|
||||
type CompactionPreparation,
|
||||
type CompactionSettings,
|
||||
calculateContextTokens,
|
||||
compact,
|
||||
|
|
@ -150,6 +151,17 @@ function createCustomMessageEntry(content: string): CustomMessageEntry {
|
|||
return entry;
|
||||
}
|
||||
|
||||
function requireSourceMessages(
|
||||
preparation: CompactionPreparation | undefined,
|
||||
): asserts preparation is CompactionPreparation & {
|
||||
sourceMessages: AgentMessage[];
|
||||
turnPrefixSourceMessages: AgentMessage[];
|
||||
} {
|
||||
if (!preparation?.sourceMessages || !preparation.turnPrefixSourceMessages) {
|
||||
throw new Error("Expected compaction source messages");
|
||||
}
|
||||
}
|
||||
|
||||
function extractText(messages: AgentMessage[]): string {
|
||||
return messages
|
||||
.map((message) => {
|
||||
|
|
@ -461,6 +473,44 @@ describe("buildSessionContext", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("prepareCompaction source messages", () => {
|
||||
it("provides the active history prefix for a normal compaction", () => {
|
||||
const u1 = createMessageEntry(createUserMessage("old user ".repeat(20)));
|
||||
const a1 = createMessageEntry(createAssistantMessage("old assistant ".repeat(20)));
|
||||
const u2 = createMessageEntry(createUserMessage("recent user"));
|
||||
const a2 = createMessageEntry(createAssistantMessage("ok"));
|
||||
const settings: CompactionSettings = { ...DEFAULT_COMPACTION_SETTINGS, keepRecentTokens: 3 };
|
||||
|
||||
const preparation = prepareCompaction([u1, a1, u2, a2], settings);
|
||||
|
||||
expect(preparation).toBeDefined();
|
||||
requireSourceMessages(preparation);
|
||||
expect(preparation.isSplitTurn).toBe(false);
|
||||
expect(preparation.sourceMessages).toEqual(preparation.messagesToSummarize);
|
||||
expect(extractText(preparation.sourceMessages)).toContain("old user");
|
||||
expect(extractText(preparation.sourceMessages)).not.toContain("recent user");
|
||||
expect(preparation.turnPrefixSourceMessages).toEqual([]);
|
||||
});
|
||||
|
||||
it("provides the active prefix through a split turn", () => {
|
||||
const user = createMessageEntry(createUserMessage("large request ".repeat(20)));
|
||||
const earlyAssistant = createMessageEntry(createAssistantMessage("early work ".repeat(20)));
|
||||
const keptAssistant = createMessageEntry(createAssistantMessage("kept"));
|
||||
const settings: CompactionSettings = { ...DEFAULT_COMPACTION_SETTINGS, keepRecentTokens: 1 };
|
||||
|
||||
const preparation = prepareCompaction([user, earlyAssistant, keptAssistant], settings);
|
||||
|
||||
expect(preparation).toBeDefined();
|
||||
requireSourceMessages(preparation);
|
||||
expect(preparation.isSplitTurn).toBe(true);
|
||||
expect(preparation.messagesToSummarize).toEqual([]);
|
||||
expect(preparation.sourceMessages).toEqual([]);
|
||||
expect(preparation.turnPrefixSourceMessages).toEqual(preparation.turnPrefixMessages);
|
||||
expect(extractText(preparation.turnPrefixSourceMessages)).toContain("large request");
|
||||
expect(extractText(preparation.turnPrefixSourceMessages)).not.toContain("kept");
|
||||
});
|
||||
});
|
||||
|
||||
describe("prepareCompaction with previous compaction", () => {
|
||||
it("should skip repeated compactions when kept messages still fit", () => {
|
||||
const u1 = createMessageEntry(createUserMessage("user msg 1 (summarized by compaction1)"));
|
||||
|
|
@ -479,6 +529,33 @@ describe("prepareCompaction with previous compaction", () => {
|
|||
expect(preparation).toBeUndefined();
|
||||
});
|
||||
|
||||
it("preserves retained older summaries in the active source prefix", () => {
|
||||
const u1 = createMessageEntry(createUserMessage("user 1"));
|
||||
const a1 = createMessageEntry(createAssistantMessage("assistant 1"));
|
||||
const u2 = createMessageEntry(createUserMessage("user 2 ".repeat(20)));
|
||||
const a2 = createMessageEntry(createAssistantMessage("assistant 2 ".repeat(20)));
|
||||
const compaction1 = createCompactionEntry("First summary", u1.id);
|
||||
const u3 = createMessageEntry(createUserMessage("user 3 ".repeat(20)));
|
||||
const a3 = createMessageEntry(createAssistantMessage("assistant 3 ".repeat(20)));
|
||||
const compaction2 = createCompactionEntry("Second summary", u2.id);
|
||||
const u4 = createMessageEntry(createUserMessage("user 4 ".repeat(20)));
|
||||
const a4 = createMessageEntry(createAssistantMessage("kept"));
|
||||
const settings: CompactionSettings = { ...DEFAULT_COMPACTION_SETTINGS, keepRecentTokens: 1 };
|
||||
const pathEntries = [u1, a1, u2, a2, compaction1, u3, a3, compaction2, u4, a4];
|
||||
|
||||
const preparation = prepareCompaction(pathEntries, settings);
|
||||
|
||||
expect(preparation).toBeDefined();
|
||||
requireSourceMessages(preparation);
|
||||
const activeMessages = buildSessionContext(pathEntries).messages;
|
||||
expect(preparation.sourceMessages).toEqual(activeMessages.slice(0, preparation.sourceMessages.length));
|
||||
expect(preparation.turnPrefixSourceMessages).toEqual(
|
||||
activeMessages.slice(0, preparation.turnPrefixSourceMessages.length),
|
||||
);
|
||||
const sourceSummaries = preparation.sourceMessages.filter((message) => message.role === "compactionSummary");
|
||||
expect(sourceSummaries.map((message) => message.summary)).toEqual(["Second summary", "First summary"]);
|
||||
});
|
||||
|
||||
it("should re-summarize previously kept messages when the recent window moves past them", () => {
|
||||
const u1 = createMessageEntry(createUserMessage("user msg 1 (summarized by compaction1)".repeat(4)));
|
||||
const a1 = createMessageEntry(createAssistantMessage("assistant msg 1".repeat(4)));
|
||||
|
|
@ -497,11 +574,15 @@ describe("prepareCompaction with previous compaction", () => {
|
|||
const preparation = prepareCompaction([u1, a1, u2, a2, u3, a3, compaction1, u4, a4], settings);
|
||||
|
||||
expect(preparation).toBeDefined();
|
||||
const summarizedText = extractText(preparation!.messagesToSummarize);
|
||||
requireSourceMessages(preparation);
|
||||
const summarizedText = extractText(preparation.messagesToSummarize);
|
||||
expect(summarizedText).toContain("user msg 2 - kept by compaction1");
|
||||
expect(summarizedText).toContain("user msg 3 - kept by compaction1");
|
||||
expect(summarizedText).not.toContain("First summary");
|
||||
expect(preparation!.previousSummary).toBe("First summary");
|
||||
expect(preparation.previousSummary).toBe("First summary");
|
||||
expect(preparation.sourceMessages[0]?.role).toBe("compactionSummary");
|
||||
expect(extractText(preparation.sourceMessages)).toContain("First summary");
|
||||
expect(preparation.sourceMessages.slice(1)).toEqual(preparation.messagesToSummarize);
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue