285 lines
11 KiB
TypeScript
285 lines
11 KiB
TypeScript
|
import type { ModelMessage } from 'ai';
|
||
|
|
import { redactSecretsForModel, redactSecretsInValueForModel } from './modelSecretRedaction';
|
||
|
|
|
||
|
|
const SUPERSEDED_READ_PREFIX = '[superseded read:';
|
||
|
|
const EARLIER_TERMINAL_PREFIX = '[earlier terminal output omitted:';
|
||
|
|
|
||
|
|
function isRecord(value: unknown): value is Record<string, unknown> {
|
||
|
|
return typeof value === 'object' && value !== null;
|
||
|
|
}
|
||
|
|
|
||
|
|
function getToolCallMap(messages: ModelMessage[]): Map<string, { toolName: string; input: unknown }> {
|
||
|
|
const map = new Map<string, { toolName: string; input: unknown }>();
|
||
|
|
for (const message of messages) {
|
||
|
|
if (message.role !== 'assistant' || !Array.isArray(message.content)) continue;
|
||
|
|
for (const part of message.content as unknown[]) {
|
||
|
|
if (!isRecord(part) || part.type !== 'tool-call') continue;
|
||
|
|
const toolCallId = typeof part.toolCallId === 'string' ? part.toolCallId : '';
|
||
|
|
const toolName = typeof part.toolName === 'string' ? part.toolName : '';
|
||
|
|
if (toolCallId) {
|
||
|
|
map.set(toolCallId, { toolName, input: part.input });
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return map;
|
||
|
|
}
|
||
|
|
|
||
|
|
function getToolResultParts(message: ModelMessage): Array<Record<string, unknown>> {
|
||
|
|
if (message.role !== 'tool' || !Array.isArray(message.content)) return [];
|
||
|
|
return (message.content as unknown[]).filter((part) => {
|
||
|
|
return isRecord(part) && part.type === 'tool-result';
|
||
|
|
}) as Array<Record<string, unknown>>;
|
||
|
|
}
|
||
|
|
|
||
|
|
function getToolResultText(part: Record<string, unknown>): string {
|
||
|
|
const output = part.output;
|
||
|
|
if (isRecord(output) && output.type === 'text' && typeof output.value === 'string') {
|
||
|
|
return output.value;
|
||
|
|
}
|
||
|
|
if (typeof output === 'string') return output;
|
||
|
|
return '';
|
||
|
|
}
|
||
|
|
|
||
|
|
function isToolResultError(part: Record<string, unknown>, text: string): boolean {
|
||
|
|
if (part.isError === true) return true;
|
||
|
|
if (text.toLowerCase().includes('error')) return true;
|
||
|
|
return isToolResultErrorOutput(part.output);
|
||
|
|
}
|
||
|
|
|
||
|
|
function isToolResultErrorOutput(output: unknown): boolean {
|
||
|
|
if (output == null) return false;
|
||
|
|
if (typeof output === 'object') {
|
||
|
|
const obj = output as Record<string, unknown>;
|
||
|
|
if ('error' in obj && typeof obj.error === 'string') return true;
|
||
|
|
if ('ok' in obj && obj.ok === false) return true;
|
||
|
|
if (obj.type === 'json' || obj.type === 'object') {
|
||
|
|
return isToolResultErrorOutput(obj.value);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if (typeof output === 'string') {
|
||
|
|
try {
|
||
|
|
const parsed = JSON.parse(output) as Record<string, unknown>;
|
||
|
|
if ('error' in parsed && typeof parsed.error === 'string') return true;
|
||
|
|
if ('ok' in parsed && parsed.ok === false) return true;
|
||
|
|
} catch {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
|
||
|
|
function isCachedOrSuperseded(text: string): boolean {
|
||
|
|
return text.includes('[cached]')
|
||
|
|
|| text.startsWith(SUPERSEDED_READ_PREFIX)
|
||
|
|
|| text.startsWith(EARLIER_TERMINAL_PREFIX);
|
||
|
|
}
|
||
|
|
|
||
|
|
function isSftpReadTool(toolName: string): boolean {
|
||
|
|
return toolName === 'sftp_read'
|
||
|
|
|| toolName === 'sftp.read'
|
||
|
|
|| toolName === 'sftp_read_file';
|
||
|
|
}
|
||
|
|
|
||
|
|
function isSafeReadOnlyToolName(toolName: string): boolean {
|
||
|
|
const segments = toolName.toLowerCase().split(/[._-]+/);
|
||
|
|
const readMarkers = new Set(['read', 'get', 'list', 'search', 'fetch', 'inspect', 'status', 'info', 'poll']);
|
||
|
|
const writeMarkers = new Set([
|
||
|
|
'write', 'set', 'update', 'create', 'delete', 'remove', 'start', 'stop', 'close',
|
||
|
|
'execute', 'exec', 'run', 'upload', 'download', 'move', 'copy', 'rename', 'kill',
|
||
|
|
]);
|
||
|
|
return segments.some(segment => readMarkers.has(segment))
|
||
|
|
&& !segments.some(segment => writeMarkers.has(segment));
|
||
|
|
}
|
||
|
|
|
||
|
|
function readFingerprint(toolName: string, args: unknown): string | null {
|
||
|
|
if (!isRecord(args)) return null;
|
||
|
|
if (toolName === 'terminal_poll' || toolName === 'terminal.poll') {
|
||
|
|
const jobId = args.jobId;
|
||
|
|
if (typeof jobId !== 'string') return null;
|
||
|
|
return `terminal-poll:${jobId}:${String(args.offset ?? 0)}`;
|
||
|
|
}
|
||
|
|
if (toolName === 'terminal_read_context' || toolName === 'terminal.read_context') {
|
||
|
|
const sessionId = args.sessionId;
|
||
|
|
if (typeof sessionId !== 'string') return null;
|
||
|
|
return [
|
||
|
|
'terminal-context',
|
||
|
|
sessionId,
|
||
|
|
String(args.range ?? 'viewport'),
|
||
|
|
String(args.startLine ?? ''),
|
||
|
|
String(args.maxLines ?? ''),
|
||
|
|
].join(':');
|
||
|
|
}
|
||
|
|
if (isSftpReadTool(toolName)) {
|
||
|
|
const path = args.path ?? args.remotePath;
|
||
|
|
if (typeof path !== 'string') return null;
|
||
|
|
const sessionId = args.sessionId;
|
||
|
|
const sessionPart = typeof sessionId === 'string' ? sessionId : '';
|
||
|
|
return `read:${sessionPart}:${path}`;
|
||
|
|
}
|
||
|
|
if (toolName === 'read_attachment' || toolName === 'harness.read_attachment') {
|
||
|
|
const id = args.attachmentId ?? args.id ?? args.filename ?? args.name;
|
||
|
|
return id != null ? `attachment:${String(id)}` : null;
|
||
|
|
}
|
||
|
|
return null;
|
||
|
|
}
|
||
|
|
|
||
|
|
function terminalFingerprint(toolName: string, args: unknown): string | null {
|
||
|
|
if (toolName !== 'terminal_execute' && toolName !== 'terminal.execute') return null;
|
||
|
|
if (!isRecord(args)) return null;
|
||
|
|
const sessionId = args.sessionId;
|
||
|
|
return typeof sessionId === 'string' ? `terminal:${sessionId}` : null;
|
||
|
|
}
|
||
|
|
|
||
|
|
function replaceToolResultText(part: Record<string, unknown>, text: string): Record<string, unknown> {
|
||
|
|
const output = part.output;
|
||
|
|
if (isRecord(output) && output.type === 'text') {
|
||
|
|
return { ...part, output: { ...output, value: text } };
|
||
|
|
}
|
||
|
|
return { ...part, output: { type: 'text', value: text } };
|
||
|
|
}
|
||
|
|
|
||
|
|
function redactToolCallInputs(message: ModelMessage): ModelMessage {
|
||
|
|
if (message.role !== 'assistant' || !Array.isArray(message.content)) return message;
|
||
|
|
let changed = false;
|
||
|
|
const content = (message.content as unknown[]).map(part => {
|
||
|
|
if (!isRecord(part) || part.type !== 'tool-call' || !('input' in part)) return part;
|
||
|
|
const redacted = redactSecretsInValueForModel(part.input);
|
||
|
|
if (JSON.stringify(redacted) === JSON.stringify(part.input)) return part;
|
||
|
|
changed = true;
|
||
|
|
return { ...part, input: redacted };
|
||
|
|
});
|
||
|
|
return changed ? ({ ...message, content } as ModelMessage) : message;
|
||
|
|
}
|
||
|
|
|
||
|
|
function compressMessageToolResults(
|
||
|
|
message: ModelMessage,
|
||
|
|
updater: (toolName: string, args: unknown, text: string, isError: boolean) => string | null,
|
||
|
|
toolCallMap: Map<string, { toolName: string; input: unknown }>,
|
||
|
|
): ModelMessage {
|
||
|
|
const parts = getToolResultParts(message);
|
||
|
|
if (parts.length === 0) return message;
|
||
|
|
|
||
|
|
let changed = false;
|
||
|
|
const nextContent = (message.content as unknown[]).map((part) => {
|
||
|
|
if (!isRecord(part) || part.type !== 'tool-result') return part;
|
||
|
|
const toolCallId = typeof part.toolCallId === 'string' ? part.toolCallId : '';
|
||
|
|
const meta = toolCallMap.get(toolCallId);
|
||
|
|
const toolName = meta?.toolName ?? (typeof part.toolName === 'string' ? part.toolName : '');
|
||
|
|
const args = meta?.input;
|
||
|
|
const text = getToolResultText(part);
|
||
|
|
const isError = isToolResultError(part, text);
|
||
|
|
if (isCachedOrSuperseded(text)) return part;
|
||
|
|
const replacement = updater(toolName, args, text, isError);
|
||
|
|
if (replacement == null || replacement === text) return part;
|
||
|
|
changed = true;
|
||
|
|
return replaceToolResultText(part, replacement);
|
||
|
|
});
|
||
|
|
|
||
|
|
return changed ? ({ ...message, content: nextContent } as ModelMessage) : message;
|
||
|
|
}
|
||
|
|
|
||
|
|
export interface PruneStaleToolContextOptions {
|
||
|
|
/** Supersede stale reads and omit older terminal output only under context budget pressure. */
|
||
|
|
underBudgetPressure?: boolean;
|
||
|
|
}
|
||
|
|
|
||
|
|
export function pruneStaleToolContext(
|
||
|
|
messages: ModelMessage[],
|
||
|
|
options: PruneStaleToolContextOptions = {},
|
||
|
|
): {
|
||
|
|
messages: ModelMessage[];
|
||
|
|
didAdjust: boolean;
|
||
|
|
} {
|
||
|
|
const toolCallMap = getToolCallMap(messages);
|
||
|
|
const latestReadByKey = new Map<string, number>();
|
||
|
|
const terminalExecutionsBySession = new Map<string, Array<{ index: number; command?: string }>>();
|
||
|
|
const underBudgetPressure = options.underBudgetPressure === true;
|
||
|
|
const userTurnsAfter = new Array<number>(messages.length).fill(0);
|
||
|
|
let laterUserTurns = 0;
|
||
|
|
for (let index = messages.length - 1; index >= 0; index -= 1) {
|
||
|
|
userTurnsAfter[index] = laterUserTurns;
|
||
|
|
if (messages[index].role === 'user') laterUserTurns += 1;
|
||
|
|
}
|
||
|
|
|
||
|
|
messages.forEach((message, index) => {
|
||
|
|
for (const part of getToolResultParts(message)) {
|
||
|
|
const toolCallId = typeof part.toolCallId === 'string' ? part.toolCallId : '';
|
||
|
|
const meta = toolCallMap.get(toolCallId);
|
||
|
|
const toolName = meta?.toolName ?? (typeof part.toolName === 'string' ? part.toolName : '');
|
||
|
|
const args = meta?.input;
|
||
|
|
const readKey = readFingerprint(toolName, args);
|
||
|
|
if (readKey) {
|
||
|
|
const text = getToolResultText(part);
|
||
|
|
const isError = isToolResultError(part, text);
|
||
|
|
if (!isError) {
|
||
|
|
latestReadByKey.set(readKey, index);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
const termKey = terminalFingerprint(toolName, args);
|
||
|
|
if (underBudgetPressure && termKey) {
|
||
|
|
const callArgs = isRecord(args) ? args : {};
|
||
|
|
const entries = terminalExecutionsBySession.get(termKey) ?? [];
|
||
|
|
entries.push({
|
||
|
|
index,
|
||
|
|
command: typeof callArgs.command === 'string' ? callArgs.command : undefined,
|
||
|
|
});
|
||
|
|
terminalExecutionsBySession.set(termKey, entries);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
});
|
||
|
|
|
||
|
|
const keepTerminalIndices = new Set<number>();
|
||
|
|
if (underBudgetPressure) {
|
||
|
|
for (const entries of terminalExecutionsBySession.values()) {
|
||
|
|
for (const entry of entries.slice(-2)) {
|
||
|
|
keepTerminalIndices.add(entry.index);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
const terminalOmitByIndex = new Map<number, string>();
|
||
|
|
if (underBudgetPressure) {
|
||
|
|
for (const entries of terminalExecutionsBySession.values()) {
|
||
|
|
for (const entry of entries) {
|
||
|
|
if (keepTerminalIndices.has(entry.index)) continue;
|
||
|
|
terminalOmitByIndex.set(
|
||
|
|
entry.index,
|
||
|
|
`${EARLIER_TERMINAL_PREFIX} command=${redactSecretsForModel(entry.command ?? 'unknown')}]`,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
let didAdjust = false;
|
||
|
|
const next = messages.map((message, index) => {
|
||
|
|
const redactedMessage = redactToolCallInputs(message);
|
||
|
|
const updated = compressMessageToolResults(redactedMessage, (toolName, args, text, isError) => {
|
||
|
|
if (isError) return null;
|
||
|
|
const readKey = readFingerprint(toolName, args);
|
||
|
|
if (underBudgetPressure && readKey) {
|
||
|
|
const latestIndex = latestReadByKey.get(readKey);
|
||
|
|
if (latestIndex != null && latestIndex !== index) {
|
||
|
|
return `${SUPERSEDED_READ_PREFIX} ${readKey}]`;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
const termKey = terminalFingerprint(toolName, args);
|
||
|
|
if (underBudgetPressure && termKey && terminalOmitByIndex.has(index)) {
|
||
|
|
return terminalOmitByIndex.get(index)!;
|
||
|
|
}
|
||
|
|
if (underBudgetPressure) {
|
||
|
|
const age = userTurnsAfter[index];
|
||
|
|
if (age > 10 && !isError && isSafeReadOnlyToolName(toolName)) {
|
||
|
|
return `[older tool result omitted: tool=${toolName || 'unknown'}, chars=${text.length}]`;
|
||
|
|
}
|
||
|
|
if (age >= 3 && text.length > 4_000) {
|
||
|
|
return `${text.slice(0, 1_500)}\n\n[... tool result shortened: ${text.length - 3_000} chars omitted ...]\n\n${text.slice(-1_500)}`;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return null;
|
||
|
|
}, toolCallMap);
|
||
|
|
if (updated !== message) didAdjust = true;
|
||
|
|
return updated;
|
||
|
|
});
|
||
|
|
|
||
|
|
return { messages: next, didAdjust };
|
||
|
|
}
|