NetMesh/infrastructure/ai/harness/staleContextPruner.ts

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 };
}