memind/context-budget.mjs

import crypto from 'node:crypto';

export const CONTEXT_BUDGET_MODES = Object.freeze(['off', 'shadow', 'active']);
export const DEFAULT_CONTEXT_BUDGET_MAX_CHARS = 12_000;

const SOURCE_PRIORITY = Object.freeze({
  user_task: 100,
  skill_prompt: 90,
  goal: 85,
  temporal: 75,
  memory: 65,
  harness: 50,
});

const HARNESS_TITLE_PRIORITY = Object.freeze({
  'TKMind 用户空间沙箱': 95,
  'TKMind 当前时间基准': 94,
  'TKMind 代码委托策略': 88,
  'TKMind 日程写入规则': 82,
  'TKMind 定时自动任务规则': 82,
  'TKMind 长期记忆规则': 58,
  'TKMind 用户偏好画像': 57,
  'TKMind 已沉淀用户记忆': 56,
});

function normalizeMode(value, fallback = 'off') {
  const raw = String(value ?? fallback).trim().toLowerCase();
  return CONTEXT_BUDGET_MODES.includes(raw) ? raw : fallback;
}

export function resolveContextBudgetMode(env = process.env) {
  return normalizeMode(env.MEMIND_CONTEXT_BUDGET_MODE, 'off');
}

export function resolveContextBudgetMaxChars(env = process.env) {
  const raw = Number(env.MEMIND_CONTEXT_BUDGET_MAX_CHARS ?? DEFAULT_CONTEXT_BUDGET_MAX_CHARS);
  return Number.isFinite(raw) && raw > 0 ? Math.floor(raw) : DEFAULT_CONTEXT_BUDGET_MAX_CHARS;
}

export function fingerprintContent(text) {
  const normalized = String(text ?? '')
    .replace(/\s+/g, ' ')
    .trim()
    .toLowerCase();
  if (!normalized) return null;
  return crypto.createHash('sha256').update(normalized).digest('hex').slice(0, 16);
}

function blockPriority(block) {
  const source = String(block?.source ?? '').trim();
  const base = SOURCE_PRIORITY[source] ?? 40;
  if (source !== 'harness') return base;
  const title = String(block?.title ?? block?.label ?? '').trim();
  return HARNESS_TITLE_PRIORITY[title] ?? base;
}

function blockChars(block) {
  return String(block?.text ?? block?.content ?? '').length;
}

function normalizeMemoryBlock(item) {
  const label = String(item?.label ?? item?.title ?? '').trim();
  const text = String(item?.text ?? item?.content ?? item?.summary ?? '').trim();
  if (!text) return null;
  return {
    source: 'memory',
    id: String(item?.id ?? item?.memoryId ?? label ?? text.slice(0, 32)),
    label: label || null,
    text,
    fingerprint: fingerprintContent(text),
    priority: SOURCE_PRIORITY.memory,
    chars: text.length,
  };
}

function normalizeHarnessBlock(entry) {
  const title = String(entry?.title ?? '').trim();
  const text = String(entry?.content ?? entry?.text ?? '').trim();
  if (!text) return null;
  return {
    source: 'harness',
    id: title || text.slice(0, 32),
    title: title || null,
    text,
    fingerprint: fingerprintContent(text),
    priority: blockPriority({ source: 'harness', title }),
    chars: text.length,
  };
}

export function buildContextBudgetBlocks({
  memories = [],
  temporalText = '',
  goalEnvelope = '',
  harnessEntries = [],
  skillPrompt = '',
  userTask = '',
} = {}) {
  const blocks = [];

  const task = String(userTask ?? '').trim();
  if (task) {
    blocks.push({
      source: 'user_task',
      id: 'user_task',
      text: task,
      fingerprint: fingerprintContent(task),
      priority: SOURCE_PRIORITY.user_task,
      chars: task.length,
    });
  }

  const skill = String(skillPrompt ?? '').trim();
  if (skill) {
    blocks.push({
      source: 'skill_prompt',
      id: 'skill_prompt',
      text: skill,
      fingerprint: fingerprintContent(skill),
      priority: SOURCE_PRIORITY.skill_prompt,
      chars: skill.length,
    });
  }

  const goal = String(goalEnvelope ?? '').trim();
  if (goal) {
    blocks.push({
      source: 'goal',
      id: 'goal',
      text: goal,
      fingerprint: fingerprintContent(goal),
      priority: SOURCE_PRIORITY.goal,
      chars: goal.length,
    });
  }

  const temporal = String(temporalText ?? '').trim();
  if (temporal) {
    blocks.push({
      source: 'temporal',
      id: 'temporal',
      text: temporal,
      fingerprint: fingerprintContent(temporal),
      priority: SOURCE_PRIORITY.temporal,
      chars: temporal.length,
    });
  }

  for (const item of Array.isArray(memories) ? memories : []) {
    const block = normalizeMemoryBlock(item);
    if (block) blocks.push(block);
  }

  for (const entry of Array.isArray(harnessEntries) ? harnessEntries : []) {
    const block = normalizeHarnessBlock(entry);
    if (block) blocks.push(block);
  }

  return blocks;
}

export function planContextBudget(input = {}, { maxChars, env = process.env } = {}) {
  const mode = resolveContextBudgetMode(env);
  const budget = Number.isFinite(maxChars) ? maxChars : resolveContextBudgetMaxChars(env);
  const blocks = buildContextBudgetBlocks(input);
  const seenFingerprints = new Set();
  const kept = [];
  const droppedItems = [];
  const duplicateItems = [];

  const sorted = [...blocks].sort((a, b) => {
    const priorityDelta = blockPriority(b) - blockPriority(a);
    if (priorityDelta !== 0) return priorityDelta;
    return blockChars(a) - blockChars(b);
  });

  let totalChars = 0;
  for (const block of sorted) {
    if (block.source === 'user_task') {
      kept.push(block);
      totalChars += block.chars;
      if (block.fingerprint) seenFingerprints.add(block.fingerprint);
      continue;
    }

    if (block.fingerprint && seenFingerprints.has(block.fingerprint)) {
      duplicateItems.push({
        source: block.source,
        id: block.id,
        reason: 'duplicate_fingerprint',
        fingerprint: block.fingerprint,
        chars: block.chars,
      });
      continue;
    }

    if (totalChars + block.chars > budget) {
      droppedItems.push({
        source: block.source,
        id: block.id,
        reason: 'over_budget',
        fingerprint: block.fingerprint ?? null,
        chars: block.chars,
        priority: blockPriority(block),
      });
      continue;
    }

    kept.push(block);
    totalChars += block.chars;
    if (block.fingerprint) seenFingerprints.add(block.fingerprint);
  }

  const fingerprints = [...seenFingerprints];
  return {
    mode,
    budget,
    blocks: kept,
    droppedItems,
    duplicateItems,
    totalChars,
    fingerprints,
    inputBlockCount: blocks.length,
    keptBlockCount: kept.length,
  };
}

export function applyContextBudgetToOrchestration({
  memoryContext = null,
  goalContext = null,
  runtimeContext = null,
  skillPrompt = '',
  userTask = '',
  plan = null,
  env = process.env,
} = {}) {
  const mode = resolveContextBudgetMode(env);
  if (mode === 'off') {
    return {
      memoryContext,
      goalContext,
      runtimeContext,
      skillPrompt,
      plan: null,
      applied: false,
    };
  }

  const resolvedPlan = plan ?? planContextBudget({
    memories: memoryContext?.injectionEnabled ? memoryContext.memories : [],
    temporalText: runtimeContext?.injectionEnabled ? runtimeContext.injectionText : '',
    goalEnvelope: goalContext?.injectionEnabled ? goalContext.envelope : '',
    skillPrompt,
    userTask,
  }, { env });

  if (mode === 'shadow') {
    return {
      memoryContext,
      goalContext,
      runtimeContext,
      skillPrompt,
      plan: resolvedPlan,
      applied: false,
    };
  }

  const keptBySource = {
    memory: new Set(),
    goal: false,
    temporal: false,
    skill_prompt: false,
  };

  for (const block of resolvedPlan.blocks) {
    if (block.source === 'memory') keptBySource.memory.add(block.id);
    if (block.source === 'goal') keptBySource.goal = true;
    if (block.source === 'temporal') keptBySource.temporal = true;
    if (block.source === 'skill_prompt') keptBySource.skill_prompt = true;
  }

  const nextMemoryContext = memoryContext?.injectionEnabled
    ? {
      ...memoryContext,
      memories: (Array.isArray(memoryContext.memories) ? memoryContext.memories : [])
        .filter((item, index) => {
          const label = String(item?.label ?? item?.title ?? '').trim();
          const text = String(item?.text ?? item?.content ?? item?.summary ?? '').trim();
          const id = String(item?.id ?? item?.memoryId ?? label ?? (text.slice(0, 32) || `memory-${index}`));
          return keptBySource.memory.has(id);
        }),
    }
    : memoryContext;

  const nextGoalContext = goalContext?.injectionEnabled && !keptBySource.goal
    ? { ...goalContext, injectionEnabled: false, envelope: '' }
    : goalContext;

  const nextRuntimeContext = runtimeContext?.injectionEnabled && !keptBySource.temporal
    ? { ...runtimeContext, injectionEnabled: false, injectionText: '' }
    : runtimeContext;

  const nextSkillPrompt = keptBySource.skill_prompt ? skillPrompt : '';

  return {
    memoryContext: nextMemoryContext,
    goalContext: nextGoalContext,
    runtimeContext: nextRuntimeContext,
    skillPrompt: nextSkillPrompt,
    plan: resolvedPlan,
    applied: true,
  };
}

export function applyContextBudgetToHarnessEntries(entries = [], { plan = null, env = process.env } = {}) {
  const mode = resolveContextBudgetMode(env);
  if (mode !== 'active') {
    return {
      entries,
      plan: plan ?? (mode === 'shadow'
        ? planContextBudget({ harnessEntries: entries }, { env })
        : null),
      applied: false,
    };
  }

  const resolvedPlan = plan ?? planContextBudget({ harnessEntries: entries }, { env });
  const keptIds = new Set(
    resolvedPlan.blocks
      .filter((block) => block.source === 'harness')
      .map((block) => block.id),
  );

  return {
    entries: entries.filter((entry) => {
      const title = String(entry?.title ?? '').trim();
      const text = String(entry?.content ?? entry?.text ?? '').trim();
      const id = title || text.slice(0, 32);
      return keptIds.has(id);
    }),
    plan: resolvedPlan,
    applied: true,
  };
}

export function buildContextBudgetResolvedEvent(plan) {
  if (!plan) return null;
  return {
    mode: plan.mode,
    budget: plan.budget,
    totalChars: plan.totalChars,
    inputBlockCount: plan.inputBlockCount,
    keptBlockCount: plan.keptBlockCount,
    droppedCount: plan.droppedItems.length,
    duplicateCount: plan.duplicateItems.length,
    droppedItems: plan.droppedItems,
    duplicateItems: plan.duplicateItems,
    fingerprints: plan.fingerprints,
  };
}