268 lines
11 KiB
TypeScript
268 lines
11 KiB
TypeScript
import type { Message, ToolDef } from '../../llm/openai-compat.js';
|
||
import type { ContextManager } from '../context-manager.js';
|
||
import type { HistorySummarizationConfig } from '../../config.js';
|
||
import { dedupeFileReads } from './file-read-dedup.js';
|
||
import { summarizeHistory } from './history-compactor.js';
|
||
import {
|
||
estimateTokensFromText,
|
||
estimateMessagesTokens,
|
||
estimateToolsTokens,
|
||
PROMPT_BASE_OVERHEAD_TOKENS,
|
||
} from './token-estimate.js';
|
||
import { logger } from '../../logger.js';
|
||
|
||
export const PROMPT_GUARD_RATIO_DEFAULT = 0.8;
|
||
/**
|
||
* Fallback used only when the LLM client's preflight error message can't be
|
||
* parsed for a safe-limit value. Conservative on purpose: kicks in only in
|
||
* degraded paths.
|
||
*/
|
||
export const PROMPT_GUARD_FALLBACK_TOKENS = 24_000;
|
||
/** Tool result messages above this size become candidates for compaction. */
|
||
export const LARGE_TOOL_RESULT_TOKENS = 8_000;
|
||
/**
|
||
* Absolute cap on the headroom reserved above the prompt when deciding the
|
||
* compact-and-continue trigger. Opencode/Crush compact at `limit - reserve`
|
||
* rather than a flat percentage; a pure ratio wastes hundreds of K of usable
|
||
* context on large-context models (a 1M model at 0.8 throws away 200K it could
|
||
* have kept). We keep the ratio as the reserve for small/medium models but cap
|
||
* the absolute reserve here so large-context models retain raw context longer.
|
||
*/
|
||
export const RESERVE_CAP_TOKENS_DEFAULT = 32_000;
|
||
|
||
/**
|
||
* Compute the prompt-size ceiling that triggers compact-and-continue.
|
||
*
|
||
* reserve = min(limit * (1 - ratio), reserveCapTokens)
|
||
* maxPromptTokens = limit - reserve
|
||
*
|
||
* - Small/medium limits (≤ reserveCapTokens / (1 - ratio)): the ratio reserve
|
||
* is below the cap, so this is identical to the old `floor(limit * ratio)`.
|
||
* - Large limits: the absolute cap binds, so the trigger moves up and the model
|
||
* uses its headroom instead of compacting early.
|
||
*/
|
||
export function computeMaxPromptTokens(
|
||
limitTokens: number,
|
||
promptGuardRatio: number,
|
||
reserveCapTokens: number = RESERVE_CAP_TOKENS_DEFAULT,
|
||
): number {
|
||
const ratioReserve = limitTokens * (1 - promptGuardRatio);
|
||
const reserve = Math.min(ratioReserve, reserveCapTokens);
|
||
return Math.max(1, Math.floor(limitTokens - reserve));
|
||
}
|
||
|
||
export type GuardResult =
|
||
| {
|
||
ok: true;
|
||
estimatedTokens: number;
|
||
compacted: boolean;
|
||
deduped: boolean;
|
||
summarized: boolean;
|
||
feedback?: string;
|
||
}
|
||
| {
|
||
ok: false;
|
||
estimatedTokens: number;
|
||
limitTokens: number;
|
||
message: string;
|
||
};
|
||
|
||
export interface GuardOptions {
|
||
promptGuardRatio?: number;
|
||
historySummarization?: HistorySummarizationConfig;
|
||
runIsolatedLlm?: (messages: Message[]) => Promise<string>;
|
||
}
|
||
|
||
export function looksLikeLargeEncodedPayload(text: string): boolean {
|
||
if (text.length < 8_000) return false;
|
||
if (/data:[^;,\s]+;base64,[A-Za-z0-9+/=\s]{2000,}/.test(text)) return true;
|
||
if (/base64[,:"'\s]+[A-Za-z0-9+/=\s]{2000,}/i.test(text)) return true;
|
||
return false;
|
||
}
|
||
|
||
export function buildPromptLimitAgentInstruction(estimatedTokens: number, maxPromptTokens: number): string {
|
||
return [
|
||
'前回のツール結果または会話履歴が大きすぎるため、一部の内容は LLM コンテキストに入れられませんでした。',
|
||
`推定 prompt サイズ: ${estimatedTokens.toLocaleString()} tokens / 安全上限: ${maxPromptTokens.toLocaleString()} tokens。`,
|
||
'全文を再読込しようとせず、必要な箇所を絞って調査を続けてください。',
|
||
'推奨行動: Read(offset/limit), Read(byte_offset/byte_length), Grep, または対象を絞った Bash で必要範囲だけ確認してください。',
|
||
'ユーザーに確認する前に、まず自分で範囲指定や検索に切り替えて続行してください。',
|
||
].join('\n');
|
||
}
|
||
|
||
export function parsePromptSafeLimitTokens(errorMessage: string): number | null {
|
||
const match = /safe limit ([\d,]+) tokens/i.exec(errorMessage);
|
||
if (!match) return null;
|
||
const parsed = Number.parseInt(match[1]!.replace(/,/g, ''), 10);
|
||
return Number.isFinite(parsed) && parsed > 0 ? parsed : null;
|
||
}
|
||
|
||
/**
|
||
* Replace `role: 'tool'` messages whose content exceeds LARGE_TOOL_RESULT_TOKENS
|
||
* with a short placeholder, in descending size order, until the prompt fits or
|
||
* candidates are exhausted. Mutates `messages` in place.
|
||
*
|
||
* Tracks the size delta directly instead of re-walking messages each iteration —
|
||
* O(candidates) vs the previous O(messages × candidates).
|
||
*/
|
||
export function compactOversizedToolResults(
|
||
messages: Message[],
|
||
fixedTokens: number, // non-message tokens: tool definitions + request base overhead
|
||
maxPromptTokens: number,
|
||
): { changed: boolean; estimatedTokens: number; omittedCount: number } {
|
||
let estimatedTokens = estimateMessagesTokens(messages) + fixedTokens;
|
||
let changed = false;
|
||
let omittedCount = 0;
|
||
if (estimatedTokens <= maxPromptTokens) return { changed, estimatedTokens, omittedCount };
|
||
|
||
const candidates = messages
|
||
.map((message, index) => ({
|
||
message,
|
||
index,
|
||
tokens: typeof message.content === 'string' ? estimateTokensFromText(message.content) : 0,
|
||
}))
|
||
.filter(({ message, tokens }) => message.role === 'tool' && tokens >= LARGE_TOOL_RESULT_TOKENS)
|
||
.sort((a, b) => b.tokens - a.tokens);
|
||
|
||
for (const candidate of candidates) {
|
||
if (estimatedTokens <= maxPromptTokens) break;
|
||
const content = typeof candidate.message.content === 'string' ? candidate.message.content : '';
|
||
const encodedHint = looksLikeLargeEncodedPayload(content)
|
||
? ' The omitted content appears to contain base64/data URLs.'
|
||
: '';
|
||
const placeholder = [
|
||
'[Tool result omitted before LLM request]',
|
||
`The previous tool result was too large to fit safely in the model context.${encodedHint}`,
|
||
'Use a narrower Read(offset/limit), Read(byte_offset/byte_length), Grep, or a targeted Bash command to inspect only the needed range.',
|
||
].join('\n');
|
||
const placeholderTokens = estimateTokensFromText(placeholder);
|
||
estimatedTokens = estimatedTokens - candidate.tokens + placeholderTokens;
|
||
candidate.message.content = placeholder;
|
||
changed = true;
|
||
omittedCount++;
|
||
}
|
||
|
||
return { changed, estimatedTokens, omittedCount };
|
||
}
|
||
|
||
/**
|
||
* Three-stage prompt-overflow defense, run before every LLM request.
|
||
*
|
||
* Stage 1 — dedupe duplicate file Reads (cheap, no LLM call, no info loss
|
||
* since the latest Read of each file is preserved).
|
||
* Stage 2 — compact oversized tool results (prune large tool messages).
|
||
* Stage 3 — anchored Markdown history summarization via runIsolatedLlm
|
||
* (Opencode-style); skipped if disabled in config or no LLM hook.
|
||
*
|
||
* Returns ok:false only when all three stages fail to bring the prompt under
|
||
* `promptGuardRatio` of the model context limit. The caller (executeMovement)
|
||
* decides whether to ABORT or force-transition.
|
||
*/
|
||
export async function guardPromptBeforeSend(
|
||
messages: Message[],
|
||
tools: ToolDef[],
|
||
contextManager?: ContextManager,
|
||
options: GuardOptions = {},
|
||
): Promise<GuardResult> {
|
||
const promptGuardRatio = options.promptGuardRatio ?? PROMPT_GUARD_RATIO_DEFAULT;
|
||
// tools is built once per movement and never mutated, so JSON.stringify it
|
||
// exactly once instead of recomputing inside every estimate call.
|
||
// The fixed part of every estimate is tools + the client's per-request base
|
||
// overhead — the guard must count exactly what the client preflight counts
|
||
// (see estimateRequestTokens), or prompts in the last-overhead-band pass
|
||
// here but are blocked at send time forever.
|
||
const fixedTokens = estimateToolsTokens(tools) + PROMPT_BASE_OVERHEAD_TOKENS;
|
||
if (!contextManager) {
|
||
return {
|
||
ok: true,
|
||
estimatedTokens: estimateMessagesTokens(messages) + fixedTokens,
|
||
compacted: false,
|
||
deduped: false,
|
||
summarized: false,
|
||
};
|
||
}
|
||
const limitTokens = contextManager.getContextLimit();
|
||
const reserveCapTokens = options.historySummarization?.reserveCapTokens ?? RESERVE_CAP_TOKENS_DEFAULT;
|
||
const maxPromptTokens = computeMaxPromptTokens(limitTokens, promptGuardRatio, reserveCapTokens);
|
||
|
||
let estimated = estimateMessagesTokens(messages) + fixedTokens;
|
||
if (estimated <= maxPromptTokens) {
|
||
return { ok: true, estimatedTokens: estimated, compacted: false, deduped: false, summarized: false };
|
||
}
|
||
|
||
const dedup = dedupeFileReads(messages);
|
||
if (dedup.changed) {
|
||
estimated = estimateMessagesTokens(messages) + fixedTokens;
|
||
logger.info(`[prompt-guard] file-read dedup replaced=${dedup.replacedCount} freedChars=${dedup.freedChars} estimated=${estimated}`);
|
||
}
|
||
if (estimated <= maxPromptTokens) {
|
||
return { ok: true, estimatedTokens: estimated, compacted: false, deduped: dedup.changed, summarized: false };
|
||
}
|
||
|
||
const compacted = compactOversizedToolResults(messages, fixedTokens, maxPromptTokens);
|
||
let summarized = false;
|
||
if (compacted.changed) {
|
||
const feedback = buildPromptLimitAgentInstruction(compacted.estimatedTokens, maxPromptTokens);
|
||
// Pop on overshoot: the feedback instruction is only valuable when it
|
||
// actually fits — otherwise we'd carry redundant guidance into stage 3.
|
||
messages.push({ role: 'user', content: feedback });
|
||
const estimatedWithFeedback = compacted.estimatedTokens + estimateTokensFromText(feedback);
|
||
if (estimatedWithFeedback <= maxPromptTokens) {
|
||
return {
|
||
ok: true,
|
||
estimatedTokens: estimatedWithFeedback,
|
||
compacted: true,
|
||
deduped: dedup.changed,
|
||
summarized: false,
|
||
feedback,
|
||
};
|
||
}
|
||
messages.pop();
|
||
}
|
||
estimated = compacted.estimatedTokens;
|
||
if (estimated <= maxPromptTokens) {
|
||
return {
|
||
ok: true,
|
||
estimatedTokens: estimated,
|
||
compacted: compacted.changed,
|
||
deduped: dedup.changed,
|
||
summarized: false,
|
||
};
|
||
}
|
||
|
||
const summarizationEnabled = options.historySummarization?.enabled !== false;
|
||
if (summarizationEnabled && options.runIsolatedLlm) {
|
||
const tailTurns = options.historySummarization?.tailTurns ?? 2;
|
||
const preserveRecentBudget = options.historySummarization?.preserveRecentBudget
|
||
?? Math.min(Math.floor(limitTokens * 0.25), 8_000);
|
||
const summary = await summarizeHistory(messages, {
|
||
tailTurns,
|
||
preserveRecentBudget,
|
||
runIsolatedLlm: options.runIsolatedLlm,
|
||
});
|
||
if (summary.summarized) {
|
||
summarized = true;
|
||
estimated = estimateMessagesTokens(messages) + fixedTokens;
|
||
logger.info(`[prompt-guard] history summarization complete freedChars=${summary.freedChars} estimated=${estimated}`);
|
||
if (estimated <= maxPromptTokens) {
|
||
return {
|
||
ok: true,
|
||
estimatedTokens: estimated,
|
||
compacted: compacted.changed,
|
||
deduped: dedup.changed,
|
||
summarized: true,
|
||
};
|
||
}
|
||
} else {
|
||
logger.warn(`[prompt-guard] history summarization skipped: ${summary.reason}`);
|
||
}
|
||
}
|
||
|
||
return {
|
||
ok: false,
|
||
estimatedTokens: estimated,
|
||
limitTokens,
|
||
message: `LLM request blocked before send: estimated prompt size ${estimated.toLocaleString()} tokens exceeds safe limit ${maxPromptTokens.toLocaleString()} tokens (${Math.round(promptGuardRatio * 100)}% of context ${limitTokens.toLocaleString()})${summarized ? ' (after dedup, compaction, and history summarization)' : ''}. Narrow the requested content with Read(offset/limit), Read(byte_offset/byte_length), Grep, or targeted Bash before continuing.`,
|
||
};
|
||
}
|