YeonGyu-Kim
a5983f1678
When over token limit after truncation, use session.revert to remove last message instead of attempting summarize (which would also fail). Skip summarize entirely when still over limit to prevent infinite loop. 🤖 Generated with assistance of [OhMyOpenCode](https://github.com/code-yeongyu/oh-my-opencode)
649 lines
19 KiB
TypeScript
649 lines
19 KiB
TypeScript
import type {
|
|
AutoCompactState,
|
|
DcpState,
|
|
RetryState,
|
|
TruncateState,
|
|
} from "./types";
|
|
import type { ExperimentalConfig } from "../../config";
|
|
import { RETRY_CONFIG, TRUNCATE_CONFIG } from "./types";
|
|
import { executeDynamicContextPruning } from "./pruning-executor";
|
|
import {
|
|
findLargestToolResult,
|
|
truncateToolResult,
|
|
truncateUntilTargetTokens,
|
|
} from "./storage";
|
|
import {
|
|
findEmptyMessages,
|
|
findEmptyMessageByIndex,
|
|
injectTextPart,
|
|
replaceEmptyTextParts,
|
|
} from "../session-recovery/storage";
|
|
import { log } from "../../shared/logger";
|
|
|
|
const PLACEHOLDER_TEXT = "[user interrupted]";
|
|
|
|
type Client = {
|
|
session: {
|
|
messages: (opts: {
|
|
path: { id: string };
|
|
query?: { directory?: string };
|
|
}) => Promise<unknown>;
|
|
summarize: (opts: {
|
|
path: { id: string };
|
|
body: { providerID: string; modelID: string };
|
|
query: { directory: string };
|
|
}) => Promise<unknown>;
|
|
revert: (opts: {
|
|
path: { id: string };
|
|
body: { messageID: string; partID?: string };
|
|
query: { directory: string };
|
|
}) => Promise<unknown>;
|
|
prompt_async: (opts: {
|
|
path: { sessionID: string };
|
|
body: { parts: Array<{ type: string; text: string }> };
|
|
query: { directory: string };
|
|
}) => Promise<unknown>;
|
|
};
|
|
tui: {
|
|
showToast: (opts: {
|
|
body: {
|
|
title: string;
|
|
message: string;
|
|
variant: string;
|
|
duration: number;
|
|
};
|
|
}) => Promise<unknown>;
|
|
};
|
|
};
|
|
|
|
function getOrCreateRetryState(
|
|
autoCompactState: AutoCompactState,
|
|
sessionID: string,
|
|
): RetryState {
|
|
let state = autoCompactState.retryStateBySession.get(sessionID);
|
|
if (!state) {
|
|
state = { attempt: 0, lastAttemptTime: 0 };
|
|
autoCompactState.retryStateBySession.set(sessionID, state);
|
|
}
|
|
return state;
|
|
}
|
|
|
|
|
|
|
|
function getOrCreateTruncateState(
|
|
autoCompactState: AutoCompactState,
|
|
sessionID: string,
|
|
): TruncateState {
|
|
let state = autoCompactState.truncateStateBySession.get(sessionID);
|
|
if (!state) {
|
|
state = { truncateAttempt: 0 };
|
|
autoCompactState.truncateStateBySession.set(sessionID, state);
|
|
}
|
|
return state;
|
|
}
|
|
|
|
function getOrCreateDcpState(
|
|
autoCompactState: AutoCompactState,
|
|
sessionID: string,
|
|
): DcpState {
|
|
let state = autoCompactState.dcpStateBySession.get(sessionID);
|
|
if (!state) {
|
|
state = { attempted: false, itemsPruned: 0 };
|
|
autoCompactState.dcpStateBySession.set(sessionID, state);
|
|
}
|
|
return state;
|
|
}
|
|
|
|
function sanitizeEmptyMessagesBeforeSummarize(sessionID: string): number {
|
|
const emptyMessageIds = findEmptyMessages(sessionID);
|
|
if (emptyMessageIds.length === 0) {
|
|
return 0;
|
|
}
|
|
|
|
let fixedCount = 0;
|
|
for (const messageID of emptyMessageIds) {
|
|
const replaced = replaceEmptyTextParts(messageID, PLACEHOLDER_TEXT);
|
|
if (replaced) {
|
|
fixedCount++;
|
|
} else {
|
|
const injected = injectTextPart(sessionID, messageID, PLACEHOLDER_TEXT);
|
|
if (injected) {
|
|
fixedCount++;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (fixedCount > 0) {
|
|
log("[auto-compact] pre-summarize sanitization fixed empty messages", {
|
|
sessionID,
|
|
fixedCount,
|
|
totalEmpty: emptyMessageIds.length,
|
|
});
|
|
}
|
|
|
|
return fixedCount;
|
|
}
|
|
|
|
function formatBytes(bytes: number): string {
|
|
if (bytes < 1024) return `${bytes}B`;
|
|
if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)}KB`;
|
|
return `${(bytes / (1024 * 1024)).toFixed(1)}MB`;
|
|
}
|
|
|
|
export async function getLastAssistant(
|
|
sessionID: string,
|
|
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
|
client: any,
|
|
directory: string,
|
|
): Promise<Record<string, unknown> | null> {
|
|
try {
|
|
const resp = await (client as Client).session.messages({
|
|
path: { id: sessionID },
|
|
query: { directory },
|
|
});
|
|
|
|
const data = (resp as { data?: unknown[] }).data;
|
|
if (!Array.isArray(data)) return null;
|
|
|
|
const reversed = [...data].reverse();
|
|
const last = reversed.find((m) => {
|
|
const msg = m as Record<string, unknown>;
|
|
const info = msg.info as Record<string, unknown> | undefined;
|
|
return info?.role === "assistant";
|
|
});
|
|
if (!last) return null;
|
|
return (last as { info?: Record<string, unknown> }).info ?? null;
|
|
} catch {
|
|
return null;
|
|
}
|
|
}
|
|
|
|
async function getLastMessageId(
|
|
sessionID: string,
|
|
client: Client,
|
|
directory: string,
|
|
): Promise<string | null> {
|
|
try {
|
|
const resp = await client.session.messages({
|
|
path: { id: sessionID },
|
|
query: { directory },
|
|
});
|
|
|
|
const data = (resp as { data?: unknown[] }).data;
|
|
if (!Array.isArray(data) || data.length === 0) return null;
|
|
|
|
const lastMsg = data[data.length - 1] as Record<string, unknown>;
|
|
const info = lastMsg.info as Record<string, unknown> | undefined;
|
|
return (info?.id as string) ?? null;
|
|
} catch {
|
|
return null;
|
|
}
|
|
}
|
|
|
|
function clearSessionState(
|
|
autoCompactState: AutoCompactState,
|
|
sessionID: string,
|
|
): void {
|
|
autoCompactState.pendingCompact.delete(sessionID);
|
|
autoCompactState.errorDataBySession.delete(sessionID);
|
|
autoCompactState.retryStateBySession.delete(sessionID);
|
|
autoCompactState.truncateStateBySession.delete(sessionID);
|
|
autoCompactState.dcpStateBySession.delete(sessionID);
|
|
autoCompactState.emptyContentAttemptBySession.delete(sessionID);
|
|
autoCompactState.compactionInProgress.delete(sessionID);
|
|
}
|
|
|
|
function getOrCreateEmptyContentAttempt(
|
|
autoCompactState: AutoCompactState,
|
|
sessionID: string,
|
|
): number {
|
|
return autoCompactState.emptyContentAttemptBySession.get(sessionID) ?? 0;
|
|
}
|
|
|
|
async function fixEmptyMessages(
|
|
sessionID: string,
|
|
autoCompactState: AutoCompactState,
|
|
client: Client,
|
|
messageIndex?: number,
|
|
): Promise<boolean> {
|
|
const attempt = getOrCreateEmptyContentAttempt(autoCompactState, sessionID);
|
|
autoCompactState.emptyContentAttemptBySession.set(sessionID, attempt + 1);
|
|
|
|
let fixed = false;
|
|
const fixedMessageIds: string[] = [];
|
|
|
|
if (messageIndex !== undefined) {
|
|
const targetMessageId = findEmptyMessageByIndex(sessionID, messageIndex);
|
|
if (targetMessageId) {
|
|
const replaced = replaceEmptyTextParts(
|
|
targetMessageId,
|
|
"[user interrupted]",
|
|
);
|
|
if (replaced) {
|
|
fixed = true;
|
|
fixedMessageIds.push(targetMessageId);
|
|
} else {
|
|
const injected = injectTextPart(
|
|
sessionID,
|
|
targetMessageId,
|
|
"[user interrupted]",
|
|
);
|
|
if (injected) {
|
|
fixed = true;
|
|
fixedMessageIds.push(targetMessageId);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if (!fixed) {
|
|
const emptyMessageIds = findEmptyMessages(sessionID);
|
|
if (emptyMessageIds.length === 0) {
|
|
await client.tui
|
|
.showToast({
|
|
body: {
|
|
title: "Empty Content Error",
|
|
message: "No empty messages found in storage. Cannot auto-recover.",
|
|
variant: "error",
|
|
duration: 5000,
|
|
},
|
|
})
|
|
.catch(() => {});
|
|
return false;
|
|
}
|
|
|
|
for (const messageID of emptyMessageIds) {
|
|
const replaced = replaceEmptyTextParts(messageID, "[user interrupted]");
|
|
if (replaced) {
|
|
fixed = true;
|
|
fixedMessageIds.push(messageID);
|
|
} else {
|
|
const injected = injectTextPart(
|
|
sessionID,
|
|
messageID,
|
|
"[user interrupted]",
|
|
);
|
|
if (injected) {
|
|
fixed = true;
|
|
fixedMessageIds.push(messageID);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if (fixed) {
|
|
await client.tui
|
|
.showToast({
|
|
body: {
|
|
title: "Session Recovery",
|
|
message: `Fixed ${fixedMessageIds.length} empty message(s). Retrying...`,
|
|
variant: "warning",
|
|
duration: 3000,
|
|
},
|
|
})
|
|
.catch(() => {});
|
|
}
|
|
|
|
return fixed;
|
|
}
|
|
|
|
export async function executeCompact(
|
|
sessionID: string,
|
|
msg: Record<string, unknown>,
|
|
autoCompactState: AutoCompactState,
|
|
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
|
client: any,
|
|
directory: string,
|
|
experimental?: ExperimentalConfig,
|
|
dcpForCompaction?: boolean,
|
|
): Promise<void> {
|
|
if (autoCompactState.compactionInProgress.has(sessionID)) {
|
|
await (client as Client).tui
|
|
.showToast({
|
|
body: {
|
|
title: "Compact In Progress",
|
|
message:
|
|
"Recovery already running. Please wait or start new session if stuck.",
|
|
variant: "warning",
|
|
duration: 5000,
|
|
},
|
|
})
|
|
.catch(() => {});
|
|
return;
|
|
}
|
|
autoCompactState.compactionInProgress.add(sessionID);
|
|
|
|
try {
|
|
const errorData = autoCompactState.errorDataBySession.get(sessionID);
|
|
const truncateState = getOrCreateTruncateState(autoCompactState, sessionID);
|
|
|
|
const isOverLimit =
|
|
errorData?.currentTokens &&
|
|
errorData?.maxTokens &&
|
|
errorData.currentTokens > errorData.maxTokens;
|
|
|
|
// PHASE 1: DCP (Dynamic Context Pruning) - prune duplicate tool calls first
|
|
const dcpState = getOrCreateDcpState(autoCompactState, sessionID);
|
|
if (dcpForCompaction !== false && !dcpState.attempted && isOverLimit) {
|
|
dcpState.attempted = true;
|
|
log("[auto-compact] PHASE 1: DCP triggered on token limit error", {
|
|
sessionID,
|
|
currentTokens: errorData.currentTokens,
|
|
maxTokens: errorData.maxTokens,
|
|
});
|
|
|
|
const dcpConfig = experimental?.dynamic_context_pruning ?? {
|
|
enabled: true,
|
|
notification: "detailed" as const,
|
|
protected_tools: [
|
|
"task",
|
|
"todowrite",
|
|
"todoread",
|
|
"lsp_rename",
|
|
"lsp_code_action_resolve",
|
|
],
|
|
};
|
|
|
|
try {
|
|
const pruningResult = await executeDynamicContextPruning(
|
|
sessionID,
|
|
dcpConfig,
|
|
client,
|
|
);
|
|
|
|
if (pruningResult.itemsPruned > 0) {
|
|
dcpState.itemsPruned = pruningResult.itemsPruned;
|
|
log("[auto-compact] DCP successful, proceeding to truncation", {
|
|
itemsPruned: pruningResult.itemsPruned,
|
|
tokensSaved: pruningResult.totalTokensSaved,
|
|
});
|
|
|
|
await (client as Client).tui
|
|
.showToast({
|
|
body: {
|
|
title: "Dynamic Context Pruning",
|
|
message: `Pruned ${pruningResult.itemsPruned} items (~${Math.round(pruningResult.totalTokensSaved / 1000)}k tokens). Proceeding to truncation...`,
|
|
variant: "success",
|
|
duration: 3000,
|
|
},
|
|
})
|
|
.catch(() => {});
|
|
// Continue to PHASE 2 (truncation) instead of summarizing immediately
|
|
} else {
|
|
log("[auto-compact] DCP did not prune any items", { sessionID });
|
|
}
|
|
} catch (error) {
|
|
log("[auto-compact] DCP failed", { error: String(error) });
|
|
}
|
|
}
|
|
|
|
// PHASE 2: Aggressive Truncation - always try when over limit (not experimental-only)
|
|
if (
|
|
isOverLimit &&
|
|
truncateState.truncateAttempt < TRUNCATE_CONFIG.maxTruncateAttempts
|
|
) {
|
|
log("[auto-compact] PHASE 2: aggressive truncation triggered", {
|
|
currentTokens: errorData.currentTokens,
|
|
maxTokens: errorData.maxTokens,
|
|
targetRatio: TRUNCATE_CONFIG.targetTokenRatio,
|
|
});
|
|
|
|
const aggressiveResult = truncateUntilTargetTokens(
|
|
sessionID,
|
|
errorData.currentTokens,
|
|
errorData.maxTokens,
|
|
TRUNCATE_CONFIG.targetTokenRatio,
|
|
TRUNCATE_CONFIG.charsPerToken,
|
|
);
|
|
|
|
if (aggressiveResult.truncatedCount > 0) {
|
|
truncateState.truncateAttempt += aggressiveResult.truncatedCount;
|
|
|
|
const toolNames = aggressiveResult.truncatedTools
|
|
.map((t) => t.toolName)
|
|
.join(", ");
|
|
const statusMsg = aggressiveResult.sufficient
|
|
? `Truncated ${aggressiveResult.truncatedCount} outputs (${formatBytes(aggressiveResult.totalBytesRemoved)})`
|
|
: `Truncated ${aggressiveResult.truncatedCount} outputs (${formatBytes(aggressiveResult.totalBytesRemoved)}) - continuing to summarize...`;
|
|
|
|
await (client as Client).tui
|
|
.showToast({
|
|
body: {
|
|
title: aggressiveResult.sufficient
|
|
? "Truncation Complete"
|
|
: "Partial Truncation",
|
|
message: `${statusMsg}: ${toolNames}`,
|
|
variant: aggressiveResult.sufficient ? "success" : "warning",
|
|
duration: 4000,
|
|
},
|
|
})
|
|
.catch(() => {});
|
|
|
|
log("[auto-compact] aggressive truncation completed", aggressiveResult);
|
|
|
|
if (aggressiveResult.sufficient) {
|
|
clearSessionState(autoCompactState, sessionID);
|
|
setTimeout(async () => {
|
|
try {
|
|
await (client as Client).session.prompt_async({
|
|
path: { sessionID },
|
|
body: { parts: [{ type: "text", text: "Continue" }] },
|
|
query: { directory },
|
|
});
|
|
} catch {}
|
|
}, 500);
|
|
return;
|
|
}
|
|
} else {
|
|
log("[auto-compact] no tool outputs found to truncate", { sessionID });
|
|
}
|
|
|
|
// PHASE 2.5: Revert fallback - if still over limit, remove last message
|
|
log("[auto-compact] PHASE 2.5: revert fallback - still over limit after truncation", {
|
|
sessionID,
|
|
currentTokens: errorData.currentTokens,
|
|
maxTokens: errorData.maxTokens,
|
|
});
|
|
|
|
const lastMessageId = await getLastMessageId(
|
|
sessionID,
|
|
client as Client,
|
|
directory,
|
|
);
|
|
|
|
if (lastMessageId) {
|
|
try {
|
|
await (client as Client).session.revert({
|
|
path: { id: sessionID },
|
|
body: { messageID: lastMessageId },
|
|
query: { directory },
|
|
});
|
|
|
|
await (client as Client).tui
|
|
.showToast({
|
|
body: {
|
|
title: "Message Reverted",
|
|
message: "Removed last message to reduce context. Retrying...",
|
|
variant: "warning",
|
|
duration: 3000,
|
|
},
|
|
})
|
|
.catch(() => {});
|
|
|
|
clearSessionState(autoCompactState, sessionID);
|
|
setTimeout(async () => {
|
|
try {
|
|
await (client as Client).session.prompt_async({
|
|
path: { sessionID },
|
|
body: { parts: [{ type: "text", text: "Continue" }] },
|
|
query: { directory },
|
|
});
|
|
} catch {}
|
|
}, 500);
|
|
return;
|
|
} catch (revertError) {
|
|
log("[auto-compact] revert failed", { error: String(revertError) });
|
|
}
|
|
}
|
|
}
|
|
|
|
// PHASE 3: Summarize - only when under limit (otherwise it will also fail)
|
|
if (isOverLimit) {
|
|
log("[auto-compact] skipping summarize - still over token limit", {
|
|
sessionID,
|
|
currentTokens: errorData?.currentTokens,
|
|
maxTokens: errorData?.maxTokens,
|
|
});
|
|
|
|
clearSessionState(autoCompactState, sessionID);
|
|
|
|
await (client as Client).tui
|
|
.showToast({
|
|
body: {
|
|
title: "Recovery Failed",
|
|
message: `Still over token limit (${errorData?.currentTokens}/${errorData?.maxTokens}). Please start a new session or manually compact.`,
|
|
variant: "error",
|
|
duration: 10000,
|
|
},
|
|
})
|
|
.catch(() => {});
|
|
return;
|
|
}
|
|
|
|
const retryState = getOrCreateRetryState(autoCompactState, sessionID);
|
|
|
|
if (errorData?.errorType?.includes("non-empty content")) {
|
|
const attempt = getOrCreateEmptyContentAttempt(
|
|
autoCompactState,
|
|
sessionID,
|
|
);
|
|
if (attempt < 3) {
|
|
const fixed = await fixEmptyMessages(
|
|
sessionID,
|
|
autoCompactState,
|
|
client as Client,
|
|
errorData.messageIndex,
|
|
);
|
|
if (fixed) {
|
|
setTimeout(() => {
|
|
executeCompact(
|
|
sessionID,
|
|
msg,
|
|
autoCompactState,
|
|
client,
|
|
directory,
|
|
experimental,
|
|
dcpForCompaction,
|
|
);
|
|
}, 500);
|
|
return;
|
|
}
|
|
} else {
|
|
await (client as Client).tui
|
|
.showToast({
|
|
body: {
|
|
title: "Recovery Failed",
|
|
message:
|
|
"Max recovery attempts (3) reached for empty content error. Please start a new session.",
|
|
variant: "error",
|
|
duration: 10000,
|
|
},
|
|
})
|
|
.catch(() => {});
|
|
return;
|
|
}
|
|
}
|
|
|
|
if (Date.now() - retryState.lastAttemptTime > 300000) {
|
|
retryState.attempt = 0;
|
|
autoCompactState.truncateStateBySession.delete(sessionID);
|
|
}
|
|
|
|
if (retryState.attempt < RETRY_CONFIG.maxAttempts) {
|
|
retryState.attempt++;
|
|
retryState.lastAttemptTime = Date.now();
|
|
|
|
const providerID = msg.providerID as string | undefined;
|
|
const modelID = msg.modelID as string | undefined;
|
|
|
|
if (providerID && modelID) {
|
|
try {
|
|
sanitizeEmptyMessagesBeforeSummarize(sessionID);
|
|
|
|
await (client as Client).tui
|
|
.showToast({
|
|
body: {
|
|
title: "Auto Compact",
|
|
message: `Summarizing session (attempt ${retryState.attempt}/${RETRY_CONFIG.maxAttempts})...`,
|
|
variant: "warning",
|
|
duration: 3000,
|
|
},
|
|
})
|
|
.catch(() => {});
|
|
|
|
await (client as Client).session.summarize({
|
|
path: { id: sessionID },
|
|
body: { providerID, modelID },
|
|
query: { directory },
|
|
});
|
|
|
|
setTimeout(async () => {
|
|
try {
|
|
await (client as Client).session.prompt_async({
|
|
path: { sessionID },
|
|
body: { parts: [{ type: "text", text: "Continue" }] },
|
|
query: { directory },
|
|
});
|
|
} catch {}
|
|
}, 500);
|
|
return;
|
|
} catch {
|
|
const delay =
|
|
RETRY_CONFIG.initialDelayMs *
|
|
Math.pow(RETRY_CONFIG.backoffFactor, retryState.attempt - 1);
|
|
const cappedDelay = Math.min(delay, RETRY_CONFIG.maxDelayMs);
|
|
|
|
setTimeout(() => {
|
|
executeCompact(
|
|
sessionID,
|
|
msg,
|
|
autoCompactState,
|
|
client,
|
|
directory,
|
|
experimental,
|
|
dcpForCompaction,
|
|
);
|
|
}, cappedDelay);
|
|
return;
|
|
}
|
|
} else {
|
|
await (client as Client).tui
|
|
.showToast({
|
|
body: {
|
|
title: "Summarize Skipped",
|
|
message: "Missing providerID or modelID.",
|
|
variant: "warning",
|
|
duration: 3000,
|
|
},
|
|
})
|
|
.catch(() => {});
|
|
}
|
|
}
|
|
|
|
clearSessionState(autoCompactState, sessionID);
|
|
|
|
await (client as Client).tui
|
|
.showToast({
|
|
body: {
|
|
title: "Auto Compact Failed",
|
|
message: "All recovery attempts failed. Please start a new session.",
|
|
variant: "error",
|
|
duration: 5000,
|
|
},
|
|
})
|
|
.catch(() => {});
|
|
} finally {
|
|
autoCompactState.compactionInProgress.delete(sessionID);
|
|
}
|
|
}
|