@animus-labs/cortex 0.2.0

package/src/compaction/microcompaction.ts

@@ -0,0 +1,709 @@
+/**
+ * Layer 1: Microcompaction (tool result trimming).
+ *
+ * Progressively reduces the footprint of old tool results in conversation
+ * history. Operates in-memory only via transformContext; never modifies
+ * the persisted conversation history (agent.state.messages).
+ *
+ * Two sub-mechanisms:
+ *   1. Insertion-time cap: Truncate large tool results at insertion time.
+ *   2. Cache-aware token-offset trimming: When the prompt cache has gone
+ *      cold and context utilization is above the trim floor, walk history
+ *      from newest to oldest accumulating token offsets. Tool results
+ *      within the hot zone stay full; beyond it, bookend size shrinks
+ *      linearly across the degradation span; beyond the span, results
+ *      become placeholder or clear based on tool category.
+ *
+ * When a persistResult callback is configured, full content is persisted
+ * to disk before any destructive trim (bookend, placeholder, clear) for
+ * non-reproducible and computational tools, and the in-context replacement
+ * includes the disk path so the agent can Read the content back.
+ *
+ * References:
+ *   - compaction-strategy.md (Layer 1: Tool Result Trimming)
+ *   - observational-memory-architecture.md (cache-aware L1)
+ *   - tool-result-persistence.md (proactive persistence at execution)
+ */
+
+import type { AgentMessage } from '../context-manager.js';
+import type { MicrocompactionConfig, ToolCategory } from '../types.js';
+import { estimateTokens } from '../token-estimator.js';
+
+// ---------------------------------------------------------------------------
+// Defaults
+// ---------------------------------------------------------------------------
+
+export const MICROCOMPACTION_DEFAULTS: MicrocompactionConfig = {
+  maxResultTokens: 50_000,
+  trimFloorRatio: 0.25,
+  hotZoneMinTokens: 16_000,
+  hotZoneRatio: 0.05,
+  bookendMaxChars: 2_000,
+  bookendMinChars: 256,
+  degradationSpanRatio: 0.40,
+  // extendedRetentionMultiplier intentionally undefined; resolved per-engine
+  // based on whether a persister is configured (1.0 with, 1.5 without).
+};
+
+/** Default extended retention multiplier when no persister is configured. */
+const EXTENDED_RETENTION_NO_PERSISTER = 1.5;
+/** Default extended retention multiplier when a persister is configured. */
+const EXTENDED_RETENTION_WITH_PERSISTER = 1.0;
+
+// ---------------------------------------------------------------------------
+// Tool category defaults
+// ---------------------------------------------------------------------------
+
+/**
+ * Default tool categories for built-in tools.
+ * Consumers can override and extend via config.toolCategories.
+ */
+const DEFAULT_TOOL_CATEGORIES: Record<string, ToolCategory> = {
+  Read: 'rereadable',
+  Glob: 'rereadable',
+  Grep: 'rereadable',
+  WebFetch: 'non-reproducible',
+  Bash: 'non-reproducible',
+  SubAgent: 'ephemeral',
+  TaskOutput: 'ephemeral',
+};
+
+// ---------------------------------------------------------------------------
+// Trim state cache
+// ---------------------------------------------------------------------------
+
+/**
+ * Trim instruction for a single tool result message.
+ * Describes how to transform the message content for the LLM context.
+ */
+export type TrimAction =
+  | { kind: 'full' }
+  | { kind: 'bookend'; headChars: number; tailChars: number; originalTokens: number }
+  | { kind: 'placeholder'; toolName: string; preview: string }
+  | { kind: 'clear' };
+
+/**
+ * Cached trim state. Keyed on history length and a coarse utilization band
+ * so the engine can replay identical trim decisions across consecutive
+ * transformContext calls when the underlying conversation hasn't changed
+ * meaningfully.
+ */
+export interface TrimState {
+  /** History length at the time the state was computed. */
+  historyLength: number;
+  /** Coarse utilization band: floor(usageRatio * 10). */
+  utilizationBand: number;
+  /** Map from conversation history index to trim action. */
+  actions: Map<number, TrimAction>;
+}
+
+// ---------------------------------------------------------------------------
+// Insertion-time cap (Tier 1)
+// ---------------------------------------------------------------------------
+
+/**
+ * Cap a tool result at insertion time. If the result exceeds maxResultTokens,
+ * truncate to head + tail bookend format using bookendMaxChars per side.
+ *
+ * This runs at insertion, NOT in transformContext. The truncated result is
+ * stored in conversation history.
+ */
+export function capToolResult(
+  content: string,
+  config: Pick<MicrocompactionConfig, 'maxResultTokens' | 'bookendMaxChars'>,
+): string {
+  const tokens = estimateTokens(content);
+  if (tokens <= config.maxResultTokens) {
+    return content;
+  }
+
+  const headSize = config.bookendMaxChars;
+  const tailSize = config.bookendMaxChars;
+
+  // Ensure we don't exceed the content length
+  if (headSize + tailSize >= content.length) {
+    return content;
+  }
+
+  const head = content.slice(0, headSize);
+  const tail = content.slice(-tailSize);
+  const trimmedTokens = tokens - estimateTokens(head) - estimateTokens(tail);
+
+  return `${head}\n\n... [~${trimmedTokens} tokens trimmed at insertion] ...\n\n${tail}`;
+}
+
+// ---------------------------------------------------------------------------
+// Content extraction helpers
+// ---------------------------------------------------------------------------
+
+/**
+ * Extract the text content from a message's content field.
+ * Handles both string content and content arrays.
+ * For content arrays, extracts text from 'text' parts and also from
+ * 'tool_result' parts that have a 'text' field (which is where
+ * pi-agent-core stores tool output).
+ */
+export function extractTextContent(message: AgentMessage): string {
+  if (typeof message.content === 'string') {
+    return message.content;
+  }
+
+  // Guard: content may be undefined/null (corrupted toolResult) or empty
+  if (!Array.isArray(message.content) || message.content.length === 0) {
+    return '';
+  }
+
+  return message.content
+    .filter(part => typeof part.text === 'string')
+    .map(part => part.text as string)
+    .join('');
+}
+
+/**
+ * Check if a message is a tool result.
+ * Pi-agent-core stores tool results as messages with content arrays
+ * containing tool_result type parts.
+ */
+export function isToolResultMessage(message: AgentMessage): boolean {
+  if (message.role === 'toolResult') {
+    return true;
+  }
+  if (!Array.isArray(message.content)) {
+    return false;
+  }
+  return message.content.some(part => part.type === 'tool_result');
+}
+
+/**
+ * Check if a message contains a tool use (tool call from the assistant).
+ */
+export function isToolUseMessage(message: AgentMessage): boolean {
+  if (message.role !== 'assistant') {
+    return false;
+  }
+  if (!Array.isArray(message.content)) {
+    return false;
+  }
+  return message.content.some(part => part.type === 'tool_use');
+}
+
+/**
+ * Extract the tool name from a tool result or tool use message.
+ * Returns null if the message is not a tool-related message.
+ */
+export function extractToolName(message: AgentMessage): string | null {
+  if (typeof message.toolName === 'string') {
+    return message.toolName;
+  }
+
+  if (!Array.isArray(message.content)) {
+    return null;
+  }
+
+  for (const part of message.content) {
+    if (part.type === 'tool_use' && typeof part['name'] === 'string') {
+      return part['name'];
+    }
+    if (part.type === 'tool_result' && typeof part['name'] === 'string') {
+      return part['name'];
+    }
+  }
+
+  return null;
+}
+
+/**
+ * Get the effective tool category for a tool name.
+ */
+export function getToolCategory(
+  toolName: string,
+  customCategories?: Record<string, ToolCategory>,
+): ToolCategory | undefined {
+  // Custom categories take precedence
+  if (customCategories?.[toolName]) {
+    return customCategories[toolName];
+  }
+  return DEFAULT_TOOL_CATEGORIES[toolName];
+}
+
+// ---------------------------------------------------------------------------
+// Token-offset trimming (cache-aware, progressive degradation)
+// ---------------------------------------------------------------------------
+
+/**
+ * Compute the effective hot zone size in tokens.
+ * `max(hotZoneMinTokens, contextWindow * hotZoneRatio)`.
+ */
+export function computeHotZone(
+  contextWindow: number,
+  config: Pick<MicrocompactionConfig, 'hotZoneMinTokens' | 'hotZoneRatio'>,
+): number {
+  return Math.max(config.hotZoneMinTokens, contextWindow * config.hotZoneRatio);
+}
+
+/**
+ * Resolve the effective extended retention multiplier. When the config value
+ * is undefined, use 1.0 if a persister is configured (content recoverable
+ * from disk) or 1.5 if not (content is truly lost on trim).
+ */
+export function resolveExtendedRetentionMultiplier(config: MicrocompactionConfig): number {
+  if (typeof config.extendedRetentionMultiplier === 'number') {
+    return config.extendedRetentionMultiplier;
+  }
+  return config.persistResult
+    ? EXTENDED_RETENTION_WITH_PERSISTER
+    : EXTENDED_RETENTION_NO_PERSISTER;
+}
+
+/**
+ * Determine the trim action for a tool result based on its token offset
+ * from the most recent message.
+ *
+ * @param message - The tool result message
+ * @param tokenOffset - Cumulative tokens between this message and the end of history
+ * @param hotZone - Hot zone size in tokens (preserved fully)
+ * @param degradationSpan - Degradation span in tokens (bookend size shrinks across this span)
+ * @param config - Microcompaction config
+ * @param extendedMultiplier - Resolved extended retention multiplier for non-reproducible tools
+ */
+export function computeAction(
+  message: AgentMessage,
+  tokenOffset: number,
+  hotZone: number,
+  degradationSpan: number,
+  config: MicrocompactionConfig,
+  extendedMultiplier: number,
+): TrimAction {
+  const toolName = extractToolName(message);
+  const category = toolName ? getToolCategory(toolName, config.toolCategories) : undefined;
+
+  // Non-reproducible tools get an extended hot zone
+  const effectiveHotZone = category === 'non-reproducible'
+    ? hotZone * extendedMultiplier
+    : hotZone;
+
+  // Within hot zone: keep full
+  if (tokenOffset < effectiveHotZone) {
+    return { kind: 'full' };
+  }
+
+  const distanceBeyondHotZone = tokenOffset - effectiveHotZone;
+  const t = degradationSpan > 0
+    ? Math.min(distanceBeyondHotZone / degradationSpan, 1.0)
+    : 1.0;
+
+  // Within degradation span: interpolated bookend
+  if (t < 1.0) {
+    const bookendChars = Math.max(
+      config.bookendMinChars,
+      Math.round(config.bookendMaxChars * (1 - t) + config.bookendMinChars * t),
+    );
+    const content = extractTextContent(message);
+    const tokens = estimateTokens(content);
+    return {
+      kind: 'bookend',
+      headChars: bookendChars,
+      tailChars: bookendChars,
+      originalTokens: tokens,
+    };
+  }
+
+  // Beyond degradation span: placeholder or clear based on tool category
+  if (category === 'ephemeral') {
+    return { kind: 'clear' };
+  }
+
+  // Default for rereadable, non-reproducible, computational, and unknown:
+  // placeholder preserves the breadcrumb of what was tried.
+  const content = extractTextContent(message);
+  const preview = content.slice(0, 80).replace(/\n/g, ' ').trim();
+  return {
+    kind: 'placeholder',
+    toolName: toolName ?? 'unknown',
+    preview,
+  };
+}
+
+/**
+ * Compute the full trim state by walking history from newest to oldest,
+ * accumulating token offsets.
+ */
+export function computeTrimState(
+  history: AgentMessage[],
+  contextWindow: number,
+  usageRatio: number,
+  config: MicrocompactionConfig,
+): TrimState {
+  const actions = new Map<number, TrimAction>();
+  const hotZone = computeHotZone(contextWindow, config);
+  const degradationSpan = contextWindow * config.degradationSpanRatio;
+  const extendedMultiplier = resolveExtendedRetentionMultiplier(config);
+
+  let tokenOffset = 0;
+
+  for (let i = history.length - 1; i >= 0; i--) {
+    const message = history[i]!;
+
+    if (isToolResultMessage(message)) {
+      const action = computeAction(
+        message,
+        tokenOffset,
+        hotZone,
+        degradationSpan,
+        config,
+        extendedMultiplier,
+      );
+      if (action.kind !== 'full') {
+        actions.set(i, action);
+      }
+    }
+
+    // Accumulate tokens regardless of message type so non-tool-result
+    // messages (assistant text, user messages) push older results further
+    // down the timeline correctly.
+    tokenOffset += estimateTokens(extractTextContent(message));
+  }
+
+  return {
+    historyLength: history.length,
+    utilizationBand: Math.floor(usageRatio * 10),
+    actions,
+  };
+}
+
+/**
+ * Apply a bookend trim to content.
+ */
+export function applyBookend(
+  content: string,
+  headChars: number,
+  tailChars: number,
+  originalTokens: number,
+): string {
+  if (headChars + tailChars >= content.length) {
+    return content;
+  }
+
+  const head = content.slice(0, headChars);
+  const tail = content.slice(-tailChars);
+  const trimmedTokens = originalTokens - estimateTokens(head) - estimateTokens(tail);
+
+  return `${head}\n\n... [~${trimmedTokens} tokens trimmed] ...\n\n${tail}`;
+}
+
+/**
+ * Apply a bookend trim with a persisted file reference. The agent can
+ * Read the path to recover the trimmed middle portion.
+ */
+export function applyBookendWithPersistence(
+  content: string,
+  headChars: number,
+  tailChars: number,
+  originalTokens: number,
+  toolName: string,
+  path: string,
+): string {
+  const totalChars = content.length;
+  const header = `[Result persisted: ${path} (${totalChars} chars, ~${originalTokens} tokens) -- ${toolName}]`;
+
+  if (headChars + tailChars >= content.length) {
+    return `${header}\n\n${content}`;
+  }
+
+  const head = content.slice(0, headChars);
+  const tail = content.slice(-tailChars);
+  const trimmedTokens = originalTokens - estimateTokens(head) - estimateTokens(tail);
+
+  return `${header}\n\n${head}\n\n... [~${trimmedTokens} tokens trimmed; full content at ${path}] ...\n\n${tail}`;
+}
+
+/**
+ * Compute the replacement text for a trim action (no persistence).
+ */
+function getTrimmedText(content: string, action: TrimAction): string {
+  if (action.kind === 'bookend') {
+    return applyBookend(content, action.headChars, action.tailChars, action.originalTokens);
+  }
+  if (action.kind === 'placeholder') {
+    return `[Tool result trimmed -- ${action.toolName}: "${action.preview}" -- see assistant response below for findings]`;
+  }
+  if (action.kind === 'clear') {
+    return '[Tool result cleared]';
+  }
+  return content;
+}
+
+/**
+ * Apply a trim action to a message, returning a new message.
+ *
+ * Preserves the content array structure for tool_result messages so that
+ * tool_use_id linkage is maintained. The Anthropic API requires every
+ * tool_use block to have a matching tool_result with the same tool_use_id;
+ * replacing the content with a plain string would break this contract.
+ */
+export function applyTrimAction(message: AgentMessage, action: TrimAction): AgentMessage {
+  if (action.kind === 'full') {
+    return message;
+  }
+
+  // If the message has a structured content array (e.g., runtime toolResult
+  // messages with text parts, or legacy tool_result parts), preserve the array
+  // structure and only replace text within each part.
+  if (Array.isArray(message.content) && message.content.length > 0) {
+    const newContent = message.content.map(part => {
+      const shouldTrim = typeof part.text === 'string' && (
+        part.type === 'tool_result' ||
+        (part.type === 'text' && isToolResultMessage(message))
+      );
+      if (!shouldTrim) {
+        return part;
+      }
+      const trimmed = getTrimmedText(part.text as string, action);
+      return { ...part, text: trimmed };
+    });
+    return { ...message, content: newContent };
+  }
+
+  // Plain string content: replace directly (backward compat for non-tool messages)
+  const content = typeof message.content === 'string' ? message.content : '';
+  const trimmed = getTrimmedText(content, action);
+  return { ...message, content: trimmed };
+}
+
+// ---------------------------------------------------------------------------
+// MicrocompactionEngine
+// ---------------------------------------------------------------------------
+
+/**
+ * Stateful engine for microcompaction. Caches trim state across consecutive
+ * calls with the same history length and utilization band so the same trim
+ * decisions are replayed when the conversation hasn't changed.
+ */
+export class MicrocompactionEngine {
+  private readonly config: MicrocompactionConfig;
+  private cachedState: TrimState | null = null;
+
+  constructor(config: Partial<MicrocompactionConfig> = {}) {
+    this.config = { ...MICROCOMPACTION_DEFAULTS, ...config };
+  }
+
+  /**
+   * Whether a persist callback is configured.
+   */
+  get hasPersistCallback(): boolean {
+    return typeof this.config.persistResult === 'function';
+  }
+
+  /**
+   * Get the current config (for testing/inspection).
+   */
+  getConfig(): MicrocompactionConfig {
+    return this.config;
+  }
+
+  /**
+   * Cap a tool result at insertion time.
+   */
+  capAtInsertion(content: string): string {
+    return capToolResult(content, this.config);
+  }
+
+  /**
+   * Apply microcompaction to conversation history messages.
+   * Called inside transformContext. Returns a new array; never modifies the input.
+   *
+   * Cache-aware gating:
+   *   - When `cacheCold` is false (cache warm), L1 is dormant and history
+   *     is returned untouched. This preserves cache hits during active use.
+   *   - When `cacheCold` is true, L1 may run if utilization is at or above
+   *     `trimFloorRatio` (default 25%).
+   *
+   * Token-offset trimming:
+   *   Walks history from newest to oldest. Tool results within the hot zone
+   *   stay full. Beyond the hot zone, bookend size shrinks linearly across
+   *   the degradation span. Beyond the span, results become placeholder or
+   *   clear based on tool category.
+   *
+   * Persistence:
+   *   When a persistResult callback is configured, full content is saved to
+   *   disk before any destructive action (bookend, placeholder, clear) for
+   *   non-reproducible and computational tools. The in-context replacement
+   *   includes the disk path so the agent can Read the content back.
+   *
+   * @param history - Conversation history messages (post-slot region)
+   * @param contextWindow - Total context window size in tokens
+   * @param currentTokens - Current estimated token count
+   * @param options.cacheCold - Whether the prompt cache has expired (or is unused)
+   * @returns Potentially trimmed conversation history
+   */
+  async apply(
+    history: AgentMessage[],
+    contextWindow: number,
+    currentTokens: number,
+    options: { cacheCold: boolean } = { cacheCold: true },
+  ): Promise<AgentMessage[]> {
+    if (contextWindow <= 0 || history.length === 0) {
+      return history;
+    }
+
+    // Gate 1: Cache awareness. When the cache is warm, L1 is dormant.
+    if (!options.cacheCold) {
+      return history;
+    }
+
+    // Gate 2: Trim floor. Below this utilization, no trimming even if cold.
+    const usageRatio = currentTokens / contextWindow;
+    if (usageRatio < this.config.trimFloorRatio) {
+      return history;
+    }
+
+    const utilizationBand = Math.floor(usageRatio * 10);
+
+    // Replay cached state when nothing meaningful has changed.
+    const canReplay = this.cachedState
+      && this.cachedState.historyLength === history.length
+      && this.cachedState.utilizationBand === utilizationBand;
+
+    if (!canReplay) {
+      this.cachedState = computeTrimState(history, contextWindow, usageRatio, this.config);
+    }
+
+    if (!this.cachedState || this.cachedState.actions.size === 0) {
+      return history;
+    }
+
+    const result: AgentMessage[] = [];
+    for (let i = 0; i < history.length; i++) {
+      const action = this.cachedState.actions.get(i);
+      if (action) {
+        const persisted = await this.maybePersistBeforeTrim(history[i]!, action, i);
+        result.push(persisted);
+      } else {
+        result.push(history[i]!);
+      }
+    }
+
+    return result;
+  }
+
+  /**
+   * If a persistResult callback is configured and the action is destructive
+   * (bookend, placeholder, or clear), persist each non-reproducible or
+   * computational tool_result part's content to disk individually and
+   * replace each part's text with a path-referencing representation.
+   *
+   * Multi-part messages (from parallel tool calls) get per-part treatment so
+   * each part's disk path maps back to its own content.
+   *
+   * Parts that don't qualify for persistence (rereadable, ephemeral, no
+   * callback, or non-tool-result parts) still get the standard trim action
+   * applied to their own text.
+   */
+  private async maybePersistBeforeTrim(
+    message: AgentMessage,
+    action: TrimAction,
+    messageIndex: number,
+  ): Promise<AgentMessage> {
+    if (action.kind === 'full') {
+      return message;
+    }
+
+    // No persister: fall back entirely to the standard per-part trim.
+    if (!this.config.persistResult) {
+      return applyTrimAction(message, action);
+    }
+
+    // String content: treat as a single unit (non-tool-result message).
+    if (!Array.isArray(message.content) || message.content.length === 0) {
+      return applyTrimAction(message, action);
+    }
+
+    // Per-part processing: each tool result text part gets its own persistence
+    // decision and its own replacement text.
+    const messageToolName = typeof message.toolName === 'string' ? message.toolName : null;
+    const newParts = await Promise.all(message.content.map(async part => {
+      const isLegacyToolResult = part.type === 'tool_result' && typeof part.text === 'string';
+      const isRuntimeToolResultText = message.role === 'toolResult' && part.type === 'text' && typeof part.text === 'string';
+      const isToolResultText = isLegacyToolResult || isRuntimeToolResultText;
+      if (!isToolResultText) {
+        return part;
+      }
+
+      const partText = part.text as string;
+      const toolName = typeof part['name'] === 'string' ? part['name'] as string : messageToolName;
+      const category = toolName ? getToolCategory(toolName, this.config.toolCategories) : undefined;
+
+      // Only persist non-reproducible and computational parts.
+      const shouldPersist = category === 'non-reproducible' || category === 'computational';
+      if (!shouldPersist) {
+        return { ...part, text: getTrimmedText(partText, action) };
+      }
+
+      let path: string;
+      try {
+        path = await this.config.persistResult!(partText, {
+          toolName: toolName ?? 'unknown',
+          messageIndex,
+          category,
+        });
+      } catch {
+        // Persist failed; fall back to standard trim for this part.
+        return { ...part, text: getTrimmedText(partText, action) };
+      }
+
+      const persistedText = formatPersistedReplacement(
+        partText,
+        action,
+        toolName ?? 'unknown',
+        path,
+      );
+      return { ...part, text: persistedText };
+    }));
+
+    return { ...message, content: newParts };
+  }
+
+  /**
+   * Reset the cached trim state. Called after Layer 2 compaction
+   * replaces conversation history.
+   */
+  resetCache(): void {
+    this.cachedState = null;
+  }
+
+  /**
+   * Get the cached state (for testing/inspection).
+   */
+  getCachedState(): TrimState | null {
+    return this.cachedState;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Persistence formatting helpers
+// ---------------------------------------------------------------------------
+
+/**
+ * Format the in-context replacement when content has been persisted to disk.
+ * - bookend: header + head + middle marker + tail
+ * - placeholder: one-line breadcrumb with disk path
+ * - clear: minimal marker with disk path
+ */
+function formatPersistedReplacement(
+  content: string,
+  action: TrimAction,
+  toolName: string,
+  path: string,
+): string {
+  if (action.kind === 'bookend') {
+    return applyBookendWithPersistence(
+      content, action.headChars, action.tailChars, action.originalTokens, toolName, path,
+    );
+  }
+  if (action.kind === 'placeholder') {
+    return `[Tool result persisted -- ${toolName}: "${action.preview}" -- use Read on ${path} for full content]`;
+  }
+  // clear
+  return `[Tool result persisted -- ${toolName} -- use Read on ${path} for full content]`;
+}
+